蓝色部分为用户实现,灰色部分为用户调用AMCT提供的API实现,用户在TensorFlow原始网络推理的代码中导入库,并在特定的位置调用相应API,即可实现压缩功能。
import amct_tensorflow as amct amct.set_logging_level(print_level='info', save_level='info')
推荐执行该步骤,以确保原始模型可以完成推理且精度正常;执行该步骤时,可以使用部分测试集,减少运行时间。
user_test_evaluate_model(evaluate_model, test_data)
train_graph = user_load_train_graph()
调用压缩图修改接口create_compressed_retrain_model,根据简易配置文件与用户原始模型对训练的图进行压缩前的图结构修改,插入通道稀疏或4选2结构化稀疏(稀疏方式二选一)以及量化感知训练的算子,生成用于组合压缩训练的图结构(对应图1中的序号1)。
record_file = './tmp/record.txt' retrain_ops = amct.create_compressed_retrain_model(graph=train_graph, config_defination=simple_cfg, outputs=user_model_outputs, record_file=record_file)
optimizer = tf.compat.v1.train.RMSPropOptimizer( ARGS.learning_rate, momentum=ARGS.momentum) train_op = optimizer.minimize(loss)
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_compressed_retrain_model,根据量化配置文件对推理的图进行压缩前的图结构修改,插入通道稀疏或4选2结构化稀疏(稀疏方式二选一)与量化感知训练的算子,用于后续模型固化与精度推理,同时生成记录稀疏信息和量化因子的record文件(对应图1中的序号4,5)。
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) #将量化因子写入record文件,说明:如用户未使能任何量化功能则无需执行此步骤,直接执行下一步 sess.run(retrain_ops[-1]) #固化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(frozen_quant_eval_pb, 'wb') as f: f.write(constant_graph.SerializeToString())
compressed_model_path = './result/user_model' amct.save_compressed_retrain_model(pb_model=trained_pb, outputs=user_model_outputs, record_file=record_file, save_path=compressed_model_path)
compressed_model = './results/user_model_compressed.pb' user_do_inference(compressed_model, test_data)