张量分解

分解约束

如果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所示，具体分解示例请参见获取更多样例。

具体流程如下：

1. 根据用户提供的TensorFlow模型文件，调用auto_decomposition生成张量分解后的模型文件。
2. 在TensorFlow训练模式的图的基础上，调用decompose_graph加载1输出的图结构改动信息(*.pkl)分解训练代码中的图，并加载1输出的分解后的模型权重文件，输出训练模式的分解图。
3. 对分解后的模型进行finetune，输出可以进行量化的模型。后续可以进行训练后量化或者量化感知训练。

图2为resnet_v2_50网络模型分解前后的示意图。

调用示例

1. 调用auto_decomposition分解已有TensorFlow模型。

   1 2 3 4 5

   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) # 执行张量分解

2. 修改已有训练代码，调用decompose_graph分解代码中的图，加载分解后的权重进行finetune。

   根据实际训练代码从如下Session方式和Estimator方式中选择一种。

   

   (*)表示用户已有的代码；...表示用户已有代码的省略，如下所列代码仅做示例，实际用户代码可能不同，请根据实际情况进行调整。

   • Session方式

     1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17

     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方式

     1 2 3 4 5 6 7 8 9 10 11 12

     from amct_tensorflow.tensor_decompose import decompose_graph save_path = "decom_path/xxx" # 保存分解后模型的路径，即中的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参数加载分解后的模型权重 # ...