使用EasyCompression进行模型压缩训练

目前，EasyCompression工具库仅支持TensorFlow1.x，不支持TensorFlow2.x。EasyCompression工具库已经集成在DSW和DLC支持的TensorFlow1.x官方镜像中，您可以结合实际需求选择在合适的开发环境中调用：

• 如果在DSW中调用，则创建DSW实例时选择的实例镜像需要为TensorFlow1.x镜像。关于如何创建DSW实例，详情请参见创建及管理DSW实例。

• 如果在DLC中调用，则创建深度学习训练任务时选择的节点镜像必须为PAI平台镜像中的TensorFlow1.x镜像。关于如何创建深度学习训练任务，详情请参见提交任务（通过控制台）。

EasyCompression工具库主要通过Compressor接口提供几种常用的模型压缩算法，例如剪枝、量化、结构化稀疏。您需要实例化Pruner、LSQQuantizer、Sparsifier等从Compressor派生的具体算法类，对算法进行初始化和设置。为了方便您在TensorFlow 1.x常用训练代码中集成压缩训练，EasyCompression提供了面向tf.train.MonitoredSession和tf.estimator.Estimator两种模式的压缩训练适配接口。基本使用流程如下图所示。如上图所示，在原有训练代码中，插入模型压缩相关代码（蓝色框标识），主要包括构造具体的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支持的参数如下表所示。

剪枝训练完成后，如果希望获得剪枝后的小模型，或需要对剪枝训练得到的模型进行微调训练，则可以参考如下典型的使用方法。

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

剪枝后的模型相比原模型会有一定加速，您也可以借助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支持的参数如下表所示。

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支持的参数如下表所示。