EasyCompression是PAI推出的面向TensorFlow模型的压缩建模训练工具库,实现了剪枝、量化及结构化稀疏等压缩训练算法,旨在帮助深度学习领域开发者方便快捷地完成模型压缩训练。本文介绍如何使用EasyCompression进行剪枝、量化及结构化稀疏训练。

使用限制

目前,EasyCompression工具库仅支持TensorFlow1.x,不支持TensorFlow2.x。EasyCompression工具库已经集成在PAI-DSW和PAI-DLC支持的TensorFlow1.x官方镜像中,您可以结合实际需求选择在合适的开发环境中调用:
  • 如果在PAI-DSW中调用,则创建PAI-DSW实例时选择的实例镜像需要为TensorFlow1.x镜像。关于如何创建PAI-DSW实例,详情请参见创建实例
  • 如果在PAI-DLC中调用,则创建深度学习训练任务时选择的节点镜像必须为PAI平台镜像中的TensorFlow1.x镜像。关于如何创建深度学习训练任务,详情请参见创建任务

操作流程

EasyCompression工具库主要通过Compressor接口提供几种常用的模型压缩算法,例如剪枝、量化、结构化稀疏。您需要实例化Pruner、LSQQuantizer、Sparsifier等从Compressor派生的具体算法类,对算法进行初始化和设置。为了方便您在TensorFlow 1.x常用训练代码中集成压缩训练,EasyCompression提供了面向tf.train.MonitoredSessiontf.estimator.Estimator两种模式的压缩训练适配接口。基本使用流程如下图所示。Compressor使用流程如上图所示,在原有训练代码中,插入模型压缩相关代码(蓝色框标识),主要包括构造具体的Compressor和修改训练流程控制。如果使用Session机制训练,需要在MonitorSession构造时引入EasyCompression提供的CompressionHook。如果使用Estimator机制训练,需要使用EasyCompression提供的train_and_evaluate函数替换原有训练调用入口。详细的压缩算法使用方式请参见如下链接:

剪枝训练

剪枝(Pruning)是指对神经网络模型的参数或计算节点进行裁剪,该过程中通常需要基于特定标准选择模型中适合裁剪的部分,渐进地进行裁剪和微调,最终获得精简且高效的模型。EasyCompression提供Pruner类实现剪枝训练算法,典型的使用方式如下所示。
from easycompression import CompressionHook, Pruner, PrunerConfig
import tensorflow as tf

# define model architecture ...

pruner_config = PrunerConfig(
    compression_ratio=0.5,
    step_interval=1000,
    num_trials=10,
)
pruner = Pruner(pruner_config)
compression_hook = CompressionHook(pruner, save_dir='output')

with tf.train.MonitoredSession(
    session_creator=tf.train.ChiefSessionCreator(),
    hooks=[compression_hook],
) as sess:
    # sess.run(init_op)
    while True:
        try:
            # sess.run(train_op)
        except tf.errors.OutOfRangeError:
            break
其中PrunerConfig用于配置剪枝训练的详细参数,Pruner实现封装具体的剪枝算法,并通过CompressionHook在后续训练循环中被自动调用。PrunerConfig支持的参数如下表所示。
参数 类型 描述 默认值
compression_ratio FLOAT 目标压缩(裁剪)比例,取值范围为[0, 1] 0
step_interval INT 每执行一次剪枝的训练步数间隔。设置合适的步数间隔,能够使得每次剪枝操作间进行充分的模型微调训练。该参数值必须大于0。 1
num_trials INT 逐步达到目标压缩率的剪枝次数。设置合适的剪枝次数,渐进地达到最终的压缩率,能够获得更加理想的压缩效果。该参数值必须大于0。 1
include_scopes LIST of STRING 指定期望进行剪枝的模型Scope,默认值为空LIST,表示对所有Scope进行剪枝。 []
exclude_scopes LIST of STRING 指定不期望进行剪枝的模型Scope,默认值为空LIST,表示对所有Scope进行剪枝。 []
prune_sort_range STRING 选择裁剪权重时的排序方式,支持以下排序方式:
  • "global":表示所有权重全局进行排序后,整体裁剪至目标压缩比例,不同Layer的裁剪比例可能不同。
  • "local":表示对每层权重进行排序后,裁剪至目标压缩比例,不同Layer的裁剪比例一致。
"global"
prunable_op_types LIST of STRING 可剪枝的计算节点类型,支持以下类型:
  • ["Conv2D", "MatMul"]
  • ["Conv2D"]
  • ["MatMul"]
["Conv2D", "MatMul"]
ema_alpha FLOAT 用于控制权重重要性累加的参数,该参数必须大于0。 0.95
剪枝训练完成后,如果希望获得剪枝后的小模型,或需要对剪枝训练得到的模型进行微调训练,则可以参考如下典型的使用方法。
from easycompression.pruning.utils import GapFinetune
ft_helper = GapFinetune(model_dir) # model_dir为对应剪枝训练完后得到的模型文件目录。

with tf.variable_scope(..., custom_getter=ft_helper.get_variable):
    # build model

# finetune model / save model
说明 上述操作已经将剪枝训练得到的权重载入模型中,您无需再次进行restore操作。
剪枝后的模型相比原模型会有一定加速,您也可以借助Blade获得进一步的加速效果,详情请参见优化TensorFlow模型

量化训练

低比特量化旨在将原始的单精度FLOAT 32分桶量化成位宽更小的定点整数,以达到节省访存开销、提升指令计算吞率的双重目的。但是该过程中很可能引入一定的精度损失。量化训练QAT(Quantization-Aware Training)是指在训练过程中考虑模型量化的需求,尽可能保证量化模型的性能效果。EasyCompression提供LSQQuantizer类实现量化训练算法,典型的使用方式如下所示。
from easycompression import CompressionHook, LSQQuantizer, QuantizerConfig, quantize_graph
import tensorflow as tf

# define model architecture ...

quantizer_config = QuantizerConfig(
    pretrain_model='model.ckpt',
    bits=8,
    step_interval=1000,
    num_trials=1,
)
quantize_graph(
    mode='train',
    model_dir='output',
    config=quantizer_config,
)

quantizer = LSQQuantizer(model_dir='output', config=quantizer_config)
compression_hook = CompressionHook(quantizer)

with tf.train.MonitoredSession(
    session_creator=tf.train.ChiefSessionCreator(),
    hooks=[compression_hook],
) as sess:
    # sess.run(init_op)
    while True:
        try:
            # sess.run(train_op)
        except tf.errors.OutOfRangeError:
            break
其中QuantizerConfig用于配置量化训练的详细参数,Quantizer实现封装具体的量化训练算法,并通过CompressionHook在后续训练循环中被自动调用。QuantizerConfig支持的参数如下表所示。
参数 类型 描述 默认值
bits INT 模型Activation的量化比特数,取值范围为[1, 8] 8
wbits INT 模型参数的量化比特数,取值范围为[1, 8] 8
signed BOOL 是否量化为有符号定点数,支持以下取值:
  • True:量化为有符号定点数。
  • False:量化为无符号定点数。
True
pretrained_model STRING 指定预训练的FP 32模型。 None
step_interval INT 每执行一次量化参数调整的训练步数间隔。设置合适的步数间隔,能够获得更加理想的量化训练结果。

该参数值必须大于0。

1000
num_trials INT 训练过程中量化参数调整的总次数。设置合适的调整次数,能够获得更加理想的量化训练效果。

该参数值必须大于0。

10
include_scopes LIST of STRING 指定期望进行量化的模型Scope,默认值为空LIST,表示对所有Scope进行量化。 []
exclude_scopes LIST of STRING 指定不期望进行量化的模型Scope,默认值为空LIST,表示对所有Scope进行量化。 []
fp32_layers LIST of STRING 指定期望保留为FP 32计算的网络层,取值为Node Name构成的列表。 []
post_quant BOOL 是否对输出张量(Output Activation)进行量化,支持以下取值:
  • True:对输出张量进行量化。
  • False:不对输出张量进行量化。
False
int8_winograd BOOL 是否执行INT8 Winograd-Aware量化训练。
  • True:执行INT8 Winograd-Aware量化训练。
  • False:不执行INT8 Winograd-Aware量化训练。
False

量化训练完成后,您按照常规方式导出模型即可。但是需要注意的是,该模型仍为FP 32模型,如果期望获得实际量化效果,后续可以根据部署需求借助Blade、MNN或TensorRT实现模型量化,此时量化后模型性能将与量化训练时的性能效果基本一致。关于如何借助Blade实现模型量化,详情请参见TensorFlow模型量化

结构化稀疏训练

稀疏化是指将模型参数中冗余的或不重要的部分剔除(或置为零),从而降低需要存储的参数量,减少需要参与计算的权重。但是直接执行稀疏化操作会对模型性能有一定的影响,因此通常需要通过重新训练或者微调训练保持模型效果。结构化稀疏训练是指在训练过程中考虑稀疏化需求,渐进式地将模型中每层参数处理为指定比例的稀疏度。此外,为保证稀疏化的加速效果,训练过程中将使得模型中的非零参数分布符合特定约束。

EasyCompression提供Sparsifier类实现结构稀疏训练算法,典型的使用方式如下所示。
from easycompression import CompressionHook, Sparsifier, SparsifierConfig
import tensorflow as tf

# define model architecture ...

sparsifier_config = SparsifierConfig(
    target_sparse_ratio=0.6,
    sparse_block_shape=(1, 16),
    step_interval=1000,
    num_trials=10,
)
sparsifier = Sparsifier(sparsifier_config)
compression_hook = CompressionHook(sparsifier, save_dir='output')

with tf.train.MonitoredSession(
    session_creator=tf.train.ChiefSessionCreator(),
    hooks=[compression_hook],
) as sess:
    # sess.run(init_op)
    while True:
        try:
            # sess.run(train_op)
        except tf.errors.OutOfRangeError:
            break
其中SparsifierConfig用于配置稀疏训练的详细参数,Sparsifier实现封装具体的稀疏算法,并通过CompressionHook在后续训练循环中被自动调用。SparsifierConfig支持的参数如下表所示。
参数 类型 描述 默认值
target_sparse_ratio FLOAT 目标稀疏度,取值范围为[0, 1] 0.6
sparse_block_shape [INT, INT] Block Sparsity中的Block Pattern,默认为[1, 1],即单个Weight为一个稀疏单位。

该参数值列表中的两个元素均必须大于0。

[1, 1]
step_interval INT 每执行一次稀疏化的训练步数间隔。设置合适的步数间隔,能够使得每次稀疏操作间进行充分的模型微调训练。

该参数值必须大于0。

1000
num_trials INT 逐步达到目标稀疏度的稀疏化次数。设置合适的稀疏化次数,渐进地达到最终的稀疏度,能够获得更加理想的压缩效果。

该参数值必须大于0。

10
include_scopes LIST of STRING 指定期望进行稀疏化的模型Scope。默认值为空LIST,表示对所有Scope进行稀疏化。 []
exclude_scopes LIST of STRING 指定不期望进行稀疏化的模型Scope。默认值为空LIST,表示对所有Scope进行稀疏化。 []
sort_range STRING 稀疏化权重时的排序方式,支持以下取值:
  • "global":表示所有权重全局进行排序后,整体处理至目标稀疏度,不同Layer的稀疏度可能不同。
  • "local":表示对每层权重进行排序后,处理至目标稀疏度,不同Layer的稀疏度一致。
"local"

结构化稀疏训练完成后,您按照常规方式导出模型即可。需要注意的是,虽然此时模型中大部分参数已经置零,但是仍为密集存储,如果按照常规方式进行模型推理,仍无法获得加速效果,需要结合高效的稀疏算子实现加速。例如借助Blade等优化工具获取模型稀疏化收益,详情请参见优化TensorFlow模型

后续步骤

对于使用EasyCompression进行压缩训练获得的模型,通常结合推理优化工具Blade能够获得进一步的加速效果,详情请参见优化TensorFlow模型