基础稀疏
昇腾模型压缩工具目前主要支持基于重训练的通道稀疏模型压缩特性,稀疏示例请参见获取更多样例>resnet_v1_50。支持通道稀疏的层以及约束如下:
支持的层类型 |
约束 |
---|---|
MatMul |
transpose_a=False, transpose_b=True/False,adjoint_a=False,adjoint_b=True/False 权重数据类型为 Float32, Float64 |
Conv2D |
权重数据类型为 Float32, Float64 |
蓝色部分为用户实现,灰色部分为用户调用昇腾模型压缩工具提供的API实现,用户在TensorFlow原始网络推理的代码中导入库,并在特定的位置调用相应API,即可实现稀疏功能。
简要流程如下:
- 用户构造训练模式的图结构 ,调用create_prune_retrain_model接口,根据稀疏配置文件对训练的图进行稀疏前的图结构修改:在图结构中插入mask稀疏算子。
- 训练模型,将训练后的参数保存为checkpoint文件。
- 用户构造推理模式的图结构 ,调用稀疏图修改接口create_prune_retrain_model,同时根据配置文件对推理的图进行稀疏前的图结构修改:在图结构中插入mask稀疏算子;同时生成记录稀疏信息的record文件。
- 创建会话,恢复训练参数,并将推理图固化为pb模型。
- 调用稀疏图保存接口save_prune_retrain_model,根据record文件以及固化模型,导出通道裁剪后的pb推理模型,同时删除插入的mask稀疏算子。
调用示例
- 如下示例标有“由用户补充处理”的步骤,需要用户根据自己的模型和数据集进行补充处理,示例中仅为示例代码。
- 调用昇腾模型压缩工具的部分,函数入参可以根据实际情况进行调整。量化感知训练基于用户的训练过程,请确保已经有基于TensorFlow环境进行训练的脚本,并且训练后的精度正常。
- 导入昇腾模型压缩工具包,设置日志级别。
import amct_tensorflow as amct amct.set_logging_level(print_level='info', save_level='info')
- (可选,由用户补充处理)创建图并读取训练好的参数,在TensorFlow环境下推理,验证环境、推理脚本是否正常。
推荐执行该步骤,以确保原始模型可以完成推理且精度正常;执行该步骤时,可以使用部分测试集,减少运行时间。
user_test_evaluate_model(evaluate_model, test_data)
- (由用户补充处理)创建训练图。
train_graph = user_load_train_graph()
- 调用昇腾模型压缩工具,执行带稀疏算子的训练流程。
- 在训练图中插入稀疏算子。
用户基于构造的训练模式的图结构(BN的is_training参数为True) ,调用create_prune_retrain_model接口(对应图1中的序号1),根据稀疏配置文件(对应图1中的序号2)对训练的图进行稀疏前的图结构修改。create_prune_retrain_model接口会在图结构中插入mask算子,达到推理时伪稀疏的效果。稀疏配置文件需要参见量化感知训练简易配置文件说明自行构造。
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后执行)。
optimizer = tf.compat.v1.train.RMSPropOptimizer( ARGS.learning_rate, momentum=ARGS.momentum) train_op = optimizer.minimize(loss)
- 训练模型。
创建会话,进行模型的训练,将训练后的参数保存为checkpoint文件,并得到稀疏算子的结果(对应图1中的序号3,4)。
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)
- 使用修改后的图,创建反向梯度。
- 在训练图中插入稀疏算子。
- (由用户补充处理)创建推理图。
test_graph = user_load_test_graph()
- 调用昇腾模型压缩工具,实现通道稀疏。
- 先在推理图中插入稀疏算子。
用户基于构建的推理模式的图结构(BN的is_training参数为False),调用稀疏图修改接口create_prune_retrain_model,同时根据配置文件对推理的图进行稀疏前的图结构修改(对应图1中的序号5),同时生成记录稀疏信息的record文件(对应图1中的序号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)。
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)。
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中的原始精度做对比,可以观察稀疏对精度的影响。
pruned_model = './results/user_model_pruned.pb' user_do_inference(pruned_model, test_data)
父主题: 通道稀疏