张量分解

分解约束

如果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。

具体流程如下：

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

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

调用示例

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

2. 修改已有训练代码，调用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参数加载分解后的模型权重
     # ...