TensorFlow等框架多张量获取的实现方案与overhead问题咨询

最优获取方案 #

你提到的几种方案里，两次session.run()、分别调用T1.eval()/T2.eval()都会重复执行计算路径，tf.print只能输出到日志无法在代码侧获取张量值，都不适用。最优方案是单次session.run()传入所有需要获取的张量列表，示例代码如下：

import tensorflow as tf

# 假设计算图已构建完成，T1、T2为目标张量
with tf.Session() as sess:
    sess.run(tf.global_variables_initializer())
    # 单次调用同时获取多个张量值
    t1_val, t2_val = sess.run([T1, T2])

多张量同时获取的开销逻辑 #

相比仅获取T2，单次run同时获取多个张量几乎不会产生高额额外开销，核心逻辑如下：

• TensorFlow执行session.run()前会先做计算图剪枝，仅保留所有目标张量依赖的前向路径，无依赖的节点会被直接跳过执行
• 你提供的计算图中T1是T2的前置依赖节点，单独获取T2时已经需要执行T1的计算逻辑，同时获取T1、T2的计算路径和单独获取T2完全一致，仅多了一步将T1的结果从设备（GPU/TPU）拷贝到主机内存的操作，开销可以忽略不计
• 两次session.run()才会产生高额开销：每次调用都要重新执行完整的依赖路径，相当于同一份计算跑了两次，还会额外产生一次调度、设备通信的overhead，性能远低于单次run取多值的方案

其他深度学习框架的对应实现方案 #

各个框架的实现逻辑本质上都遵循「一次前向执行返回多输出」的思路，避免重复计算：

• Caffe：完成一次前向传播后，直接通过blob_at(index)或blob_by_name(name)获取对应层的输出Blob值即可，不需要重复执行前向，一次前向可获取任意多中间层结果
• Torch（Lua版本）/PyTorch：Torch静态图模式下，单次forward执行后直接访问中间节点的.output属性即可；PyTorch默认Eager模式不需要特殊处理，前向执行过程中直接留存中间张量即可，几乎无额外开销
• CNTK：通过eval()方法传入多个需要获取的输出节点，单次调用即可返回所有节点值，底层同样做了计算图剪枝优化，不会重复执行依赖算子
• Theano：编译计算函数时将需要的多个中间张量设为输出参数，编译后的函数单次调用即可返回所有输出，仅执行一次前向计算
• Chainer：默认Eager执行模式下直接留存中间变量即可；静态图模式下同样支持单次前向调用获取多节点输出
• DyNet：动态图模式下前向计算时直接保存中间张量，静态图编译时指定多个输出即可单次调用获取，额外开销极低

框架扩展算子专属指标返回的实现思路 #

如果要为框架扩展算子专属指标返回能力，用于训练过程监控，可按以下路径实现：

1. 在算子定义层增加可配置的指标输出开关，算子执行时除了返回正常计算结果，同时按照配置生成对应指标（比如算子执行耗时、输入输出张量的均值/方差、稀疏度等）
2. 在框架的会话/执行器层增加独立的指标缓冲区，算子生成的指标直接写入缓冲区，不需要和正常张量输出走同一条返回路径，避免影响主计算流性能
3. 对外暴露统一的指标获取接口，用户可在训练迭代间隙批量拉取所有已开启的算子指标，不需要修改原有前向/反向调用逻辑
4. 支持指标自定义注册，用户可自行编写算子对应的指标计算逻辑，注册到对应算子类型上，不需要修改框架核心代码

内容的提问来源于stack exchange，提问作者f.k