张量分解
张量分解通过分解卷积核的张量,将一个卷积转化成两个小卷积的堆叠来降低推理开销,如用户模型中存在大量卷积,且卷积核shape普遍大于(64, 64, 3, 3)时推荐使用张量分解。目前仅支持满足如下条件的卷积进行分解:
- group=1,dilation=(1,1),stride<3
- kernel_h>2,kernel_w>2
卷积是否分解由AMCT自动判断,即便满足上述条件,也不一定会分解。
例如,用户使用的TensorFlow原始模型中存在Conv2D层,并且该层满足上述条件,才有可能将相应的Conv2D层分解成两个Conv2D层;然后使用AMCT转换成可以在昇腾AI处理器部署的量化模型,以便在模型推理时获得更好的性能。
该场景为可选操作,用户自行决定是否进行原始模型的分解。
分解约束
如果Conv2D层的shape过大,会造成分解时间过长或分解异常中止,为防止出现该情况,执行分解动作前,请先参见如下约束或参考数据:
- 分解工具性能参考数据:
- CPU: Intel(R) Xeon(R) CPU E5-2699 v4 @ 2.20GHz
- 内存: 512G
分解单层卷积:
- shape(512, 512, 5, 5),大约耗时25秒。
- shape(1024, 1024, 3, 3),大约耗时16秒。
- shape(1024, 1024, 5, 5),大约耗时78秒。
- shape(2048, 2048, 3, 3),大约耗时63秒。
- shape(2048, 2048, 5, 5),大约耗时430秒。
- 内存超限风险提醒:
分解大卷积核存在内存超限风险参考数值:shape为(2048, 2048, 5, 5)的卷积核约占用32G内存。
接口调用流程
接口调用流程如图1所示,具体分解示例请参见获取更多样例>tensor_decompose。
具体流程如下:
- 根据用户提供的TensorFlow模型文件,调用auto_decomposition生成张量分解后的模型文件。
- 在TensorFlow训练模式的图的基础上,调用decompose_graph加载1输出的图结构改动信息(*.pkl)分解训练代码中的图,并加载1输出的分解后的模型权重文件,输出训练模式的分解图。
- 对分解后的模型进行finetune,输出可以进行量化的模型。后续可以进行训练后量化或者量化感知训练。
图2为resnet_v2_50网络模型分解前后的示意图。
调用示例
- 调用auto_decomposition分解已有TensorFlow模型。
from amct_tensorflow.tensor_decompose import auto_decomposition meta_path = "src_path/xxx.meta" # meta文件路径 ckpt_path = "src_path/xxx" # ckpt文件路径,例如,原始模型文件为src_path下的xxx.data-XXXXX-of-XXXXX和xxx.index,则设该路径为src_path/xxx save_path = "decomposed_path/xxx" # 保存分解后模型的路径,如设为decomposed_path/xxx,则会在decomposed_path下生成xxx.data-XXXXX-of-XXXXX等文件 auto_decomposition(meta_path, ckpt_path, save_path) # 执行张量分解
- 修改已有训练代码,调用decompose_graph分解代码中的图,加载分解后的权重进行finetune。
根据实际训练代码从如下Session方式和Estimator方式中选择一种。
(*)表示用户已有的代码;...表示用户已有代码的省略,如下所列代码仅做示例,实际用户代码可能不同,请根据实际情况进行调整。
- Session方式
from amct_tensorflow.tensor_decompose import decompose_graph save_path = "decom_path/xxx" # 保存分解后模型的路径。即第1步中的save_path。 # ... net_output = build_net(net_input, ...) # (*)构建网络的图 decompose_graph(save_path) # 对构建出的图进行分解,请注意该步骤必须插入在构建网络图之后、应用优化器之前 variables_to_restore = tf.global_variables() # 网络图中的所有变量 restorer = tf.train.Saver(variables_to_restore) # 设置网络图中的变量为要加载的变量 loss = build_loss(net_output, ...) # (*)构造loss optimizer = build_optimizer(...) # (*)构造优化器 train_op = optimizer.minimize(loss, ...) # (*)应用优化器优化loss # ... variables_to_init = [v for v in tf.global_variables() if v not in variables_to_restore] # 不会加载的变量 init = tf.variables_initializer(variables_to_init) # 为不会加载的变量准备初始化 with tf.Session() as sess: # (*)训练用Session sess.run(init) # 初始化不会加载的变量 restorer.restore(sess, save_path) # 从分解后的模型权重中加载要加载的变量 # ...
- Estimator方式
from amct_tensorflow.tensor_decompose import decompose_graph save_path = "decom_path/xxx" # 保存分解后模型的路径,即1中的save_path # ... def model_fn(features, labels, ...): # (*)Estimator的模型函数 net_output = build_net(net_input, ...) # (*)构建网络的图 decompose_graph(save_path) # 对构建出的图进行分解,请注意该步骤必须插入在构建网络图之后、应用优化器之前 loss = build_loss(net_output, ...) # (*)构造loss optimizer = build_optimizer(...) # (*)构造优化器 train_op = optimizer.minimize(loss, ...) # (*)应用优化器优化loss # ... estimator = tf.estimator.Estimator(model_fn, warm_start_from=save_path, ...) # (*)构造Estimator,通过warm_start_from参数加载分解后的模型权重 # ...
- Session方式
父主题: AMCT工具(TensorFlow)