基础稀疏
AMCT目前主要支持基于重训练的通道稀疏模型压缩特性,稀疏示例请参见获取更多样例>resnet_v1_50。支持通道稀疏的层以及约束如下:
支持的层类型 |
约束 |
---|---|
MatMul |
transpose_a=False, transpose_b=True/False,adjoint_a=False,adjoint_b=True/False 权重数据类型为 Float32, Float64 |
Conv2D |
|
蓝色部分为用户实现,灰色部分为用户调用AMCT提供的API实现,用户在TensorFlow原始网络推理的代码中导入库,并在特定的位置调用相应API,即可实现稀疏功能。
简要流程如下:
- 用户构造训练模式的图结构 ,调用create_prune_retrain_model接口,根据稀疏配置文件对训练的图进行稀疏前的图结构修改:在图结构中插入mask稀疏算子。
- 训练模型,将训练后的参数保存为checkpoint文件。
- 用户构造推理模式的图结构 ,调用稀疏图修改接口create_prune_retrain_model,同时根据配置文件对推理的图进行稀疏前的图结构修改:在图结构中插入mask稀疏算子;同时生成记录稀疏信息的record文件。
- 创建会话,恢复训练参数,并将推理图固化为pb模型。
- 调用稀疏图保存接口save_prune_retrain_model,根据record文件以及固化模型,导出通道裁剪后的pb推理模型,同时删除插入的mask稀疏算子。
调用示例
- 如下示例标有“由用户补充处理”的步骤,需要用户根据自己的模型和数据集进行补充处理,示例中仅为示例代码。
- 调用AMCT的部分,函数入参可以根据实际情况进行调整。量化感知训练基于用户的训练过程,请确保已经有基于TensorFlow环境进行训练的脚本,并且训练后的精度正常。
- 导入AMCT包,设置日志级别。
1 2
import amct_tensorflow as amct amct.set_logging_level(print_level='info', save_level='info')
- (可选,由用户补充处理)创建图并读取训练好的参数,在TensorFlow环境下推理,验证环境、推理脚本是否正常。
推荐执行该步骤,以确保原始模型可以完成推理且精度正常;执行该步骤时,可以使用部分测试集,减少运行时间。
1
user_test_evaluate_model(evaluate_model, test_data)
- (由用户补充处理)创建训练图。
1
train_graph = user_load_train_graph()
- 调用AMCT,执行带稀疏算子的训练流程。
- 在训练图中插入稀疏算子。
用户基于构造的训练模式的图结构(BN的is_training参数为True) ,调用create_prune_retrain_model接口(对应图1中的序号1),根据稀疏配置文件(对应图1中的序号2)对训练的图进行稀疏前的图结构修改。create_prune_retrain_model接口会在图结构中插入mask算子,达到推理时伪稀疏的效果。稀疏配置文件需要参见量化感知训练简易配置文件自行构造。
1 2 3 4 5 6
record_file = './tmp/record.txt' simple_cfg = './retrain.cfg' amct.create_prune_retrain_model(graph=train_graph, outputs=user_model_outputs, record_file=record_file, config_defination=simple_cfg)
由于tf.graph不具有直接删减通道的能力,修改网络中某一个算子的规格,TensorFlow会校验图结构的正确性,导致无法删减通道。因此TensorFlow通过增加mask算子,训练时通过mask算子来屏蔽掉被稀疏通道对整网训练的影响。
- (由用户补充处理)使用修改后的图,创建反向梯度,在训练集上做训练。
- 使用修改后的图,创建反向梯度。
调用自适应学习率优化器(比如RMSPropOptimizer)建立反向梯度图(该步骤需要在1后执行)。
1 2 3
optimizer = tf.compat.v1.train.RMSPropOptimizer( ARGS.learning_rate, momentum=ARGS.momentum) train_op = optimizer.minimize(loss)
- 训练模型。
创建会话,进行模型的训练,将训练后的参数保存为checkpoint文件,并得到稀疏算子的结果(对应图1中的序号3,4)。
1 2 3 4 5
with tf.Session() as sess: sess.run(tf.compat.v1.global_variables_initializer()) sess.run(outputs) #将训练后的参数保存为checkpoint文件 saver_save.save(sess, retrain_ckpt, global_step=0)
- 使用修改后的图,创建反向梯度。
- 在训练图中插入稀疏算子。
- (由用户补充处理)创建推理图。
1
test_graph = user_load_test_graph()
- 调用AMCT,实现通道稀疏。
- 先在推理图中插入稀疏算子。
用户基于构建的推理模式的图结构(BN的is_training参数为False),调用稀疏图修改接口create_prune_retrain_model,同时根据配置文件对推理的图进行稀疏前的图结构修改(对应图1中的序号5),同时生成记录稀疏信息的record文件(对应图1中的序号6)。
1 2 3 4 5 6
record_file = './tmp/record.txt' simple_cfg = './retrain.cfg' amct.create_prune_retrain_model(graph=test_graph, outputs=user_model_outputs, record_file=record_file, config_defination=simple_cfg)
- 恢复4中训练得到的checkpoint权重并固化为pb模型。
创建会话,恢复训练参数,并将推理图固化为pb模型(对应图1中的序号7, 8)。
1 2 3 4 5 6 7 8 9 10 11
variables_to_restore = tf.compat.v1.global_variables() saver_restore = tf.compat.v1.train.Saver(variables_to_restore) with tf.Session() as sess: sess.run(tf.compat.v1.global_variables_initializer()) #恢复训练参数 saver_restore.restore(sess, retrain_ckpt) #固化pb模型 constant_graph = tf.compat.v1.graph_util.convert_variables_to_constants( sess, eval_graph.as_graph_def(), [output.name[:-2] for output in outputs]) with tf.io.gfile.GFile(masked_pb_path, 'wb') as f: f.write(constant_graph.SerializeToString())
- 保存模型,删除插入的mask稀疏算子,并实现通道稀疏。
根据record文件以及固化模型,导出通道裁剪后的pb推理模型(对应图1中的序号6,9)。
1 2 3 4 5
pruned_model_path = './result/user_model' amct.save_prune_retrain_model(pb_model=masked_pb_path, outputs=user_model_outputs, record_file=record_file, save_path=pruned_model_path)
- 先在推理图中插入稀疏算子。
- (可选,由用户补充处理)使用稀疏后模型user_model_pruned.pb和测试集,在TensorFlow的环境下推理,测试稀疏后的仿真模型精度。使用稀疏后仿真模型精度与2中的原始精度做对比,可以观察稀疏对精度的影响。
1 2
pruned_model = './results/user_model_pruned.pb' user_do_inference(pruned_model, test_data)
父主题: 通道稀疏