如何查看Parameter Server与Worker间Send/Recv节点的张量值？

刚好之前调试分布式TensorFlow的时候折腾过类似的需求，我来给你详细拆解下：

首先直接给结论：默认情况下，TensorFlow的日志系统不会记录Send/Recv节点的张量值——这类细粒度的张量日志会带来极大的性能开销，官方默认是关闭的。如果要实现这个需求，你需要修改TensorFlow的底层源码，具体位置和步骤如下：

一、修改Send/Recv Op的核心实现 #

Send和Recv操作的核心逻辑在TensorFlow的C++内核代码里，具体路径是 tensorflow/core/kernels/sendrecv_ops.cc：

1. 记录Send节点的张量值 #

找到SendOp类的Compute方法，这个方法负责把张量发送到参数服务器或其他Worker节点。在方法里，当获取到要发送的张量后，添加日志打印代码：

void Compute(OpKernelContext* context) override {
  // 原有代码：获取要发送的张量及相关信息
  const Tensor& val = context->input(0);
  const string& send_device = context->input(1).scalar<tstring>()();
  const string& recv_device = context->input(2).scalar<tstring>()();
  const string& tensor_name = context->input(3).scalar<tstring>()();

  // 新增：打印张量的基本信息和具体值
  VLOG(1) << "SendOp: 发送张量 [" << tensor_name << "] 从设备 " 
          << send_device << " 到设备 " << recv_device << "\n" 
          << "张量内容:\n" << val.DebugString();

  // 原有发送逻辑（保持不变）...
}

这里用VLOG(1)是指定日志级别为1，启动TensorFlow时需要加上--v=1参数才能看到这些调试日志。

2. 记录Recv节点的张量值 #

同样在sendrecv_ops.cc里找到RecvOp类的Compute方法，在成功接收到张量后添加打印代码：

void Compute(OpKernelContext* context) override {
  // 原有接收逻辑（保持不变）...
  Tensor* val = context->mutable_output(0);
  OP_REQUIRES_OK(context, context->allocate_output(0, shape, val));

  // 新增：打印接收到的张量信息
  const string& recv_device = context->input(0).scalar<tstring>()();
  const string& send_device = context->input(1).scalar<tstring>()();
  const string& tensor_name = context->input(2).scalar<tstring>()();

  VLOG(1) << "RecvOp: 接收张量 [" << tensor_name << "] 从设备 " 
          << send_device << " 到设备 " << recv_device << "\n" 
          << "张量内容:\n" << val->DebugString();

  // 原有后续逻辑（保持不变）...
}

二、关键注意事项 #

• 性能影响：打印张量值会显著拖慢分布式训练的速度，只适合调试阶段使用，生产环境一定要把这些代码删掉。
• 日志收集：日志会输出在对应的Worker或Parameter Server节点的本地日志文件里，你需要确保能收集到所有节点的日志才能完整查看交互过程。
• 源码编译：修改完代码后，你需要重新从源码编译TensorFlow，替换掉原来的安装包才能让修改生效。

替代方案（可选） #

如果你不想折腾源码修改，也可以试试TensorFlow自带的调试工具tf.debugging.experimental.enable_dump_debug_info，不过这个工具默认不会记录Send/Recv的张量，需要配置更细的参数，效果可能不如直接改源码来得直观。

内容的提问来源于stack exchange，提问作者sfv