如何可视化实际的基础计算代码？

Question

我正在研究Tensorflow。

在一些测试中，我想查看Tensorflow计算的实际代码。 实际上，在定义图形后，它将进行将要计算的会话。

但是，我无法弄清楚实际计算（例如a * b，a + b，e ^ a等）的位置。 我期望代码在C ++中实现。 例如，我想查看tanh操作的代码，但是即使查看了cwise_op_tanh.cc文件，也找不到它。

我可以得到一些建议吗？

Answer 1

有关添加新op（及相关）的文档中说明了操作架构。 它很灵活，可以自动生成代码，因此一开始就令人困惑。

寻找实现时，通常会在源代码中搜索REGISTER_OP字符串。 该宏向TensorFlow引擎注册了一个操作（有关详细信息，请参见上述链接-涉及更多）。 该宏应让所有“指针”知道在哪里查看。 如果每种设备类型（CPU，GPU，特殊GPU）或分支/组成都有实现，则可能会很复杂。

例如，在tanh ，我寻找tanh ，或者在源代码中似乎约定为Tanh ：

> cd /path/to/tensorflow/source
> grep -R REGISTER_OP tensorflow | grep -i tanh
tensorflow/core/ops/math_grad.cc:REGISTER_OP_GRADIENT("Tanh", TanhGrad);
tensorflow/core/ops/math_grad.cc:REGISTER_OP_GRADIENT("Atanh", AtanhGrad);
tensorflow/core/ops/math_ops.cc:REGISTER_OP("Tanh").UNARY_COMPLEX();
tensorflow/core/ops/math_ops.cc:REGISTER_OP("Atanh").UNARY_COMPLEX();
tensorflow/core/ops/math_ops.cc:REGISTER_OP("TanhGrad").UNARY_GRADIENT_COMPLEX();
tensorflow/core/ops/nn_ops.cc:REGISTER_OP("_MklTanh")
tensorflow/core/ops/nn_ops.cc:REGISTER_OP("_MklTanhGrad")

结果： tensorflow/core/ops/math_ops.cc:REGISTER_OP("Tanh").UNARY_COMPLEX(); 看起来像我们想要的。 这导致了UNARY_COMPLEX宏定义：

#define UNARY_COMPLEX()                                                  \
  Input("x: T")                                                          \
      .Output("y: T")                                                    \
      .Attr("T: {half, bfloat16, float, double, complex64, complex128}") \
      .SetShapeFn(shape_inference::UnchangedShape)

实际的实现通常是包装在SetShapeFn调用中的函数。 但是通常也像下面这样委托实现。 我们知道tanh是按系数进行运算的，经过更多搜索， tensorflow/core/kernels/cwise_op_tanh.cc和tensorflow/core/kernels/cwise_op_gpu_tanh.cu.cc是唯一注册tanh内核。 但是这一次，这些文件使用REGISTER5类的“实用程序”进行更紧凑的样板注册。 例如，我们可以找到：

REGISTER5(UnaryOp, CPU, "Tanh", functor::tanh, float, Eigen::half, double,
      complex64, complex128);

这最终调用REGISTER_KERNEL_BUILDER ---该合作伙伴的REGISTER_OP ---以注册登记的操作实际执行。 查看代码，实现是functor::tanh ，它是实际实现的包装，通常委托给Eigen线性代数库。 Eigen有时可能会自己进一步授权一些操作。

笔记：

• 一种可能加快搜索速度的方法是使用REGISTER而不是REGISTER_OP作为搜索模式。 不过，那将返回更多结果。
• 所有代码摘录均来自当前的TensorFlow版本（提交ac8e67399d75edce6a9f94afaa2adb577035966e ）。 当心代码经常更改。
• 由于种种原因，Tensorflow正在迁移到更抽象的架构XLA 。 因此，以上解释将变得过时。