Question

我正在学习Tensorflow。

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

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

我能得到一些建议吗？

Answer 1

documentation on adding a new op (and related)解释了操作架构。它对于自动代码生成非常灵活，因此最初会让人感到困惑。

在寻找实现时，我通常会在源代码中搜索REGISTER_OP字符串。此宏向TensorFlow引擎注册一个操作（有关详细信息，请参阅前面提到的链接 - 它涉及更多）。宏应该给出所有＆＃34;指针＆＃34;知道在哪里看。当每种设备类型（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的内核。但是这次这些文件使用＆＃34; utilities＆＃34;像REGISTER5更为紧凑的锅炉板注册。例如，我们可以找到：

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

最终调用REGISTER_KERNEL_BUILDER --- REGISTER_OP partner ---来注册注册操作的实际实现。查看代码，实现是functor::tanh，它是实际实现的包装器，通常委托给Eigen线性代数库。而Eigen有时可能会进一步delegate一些操作。

注意：

可能加快搜索速度的一种方法是使用REGISTER代替REGISTER_OP作为搜索模式。不过，这会带来更多结果。
所有代码提取都来自当前的TensorFlow版本（commit ac8e67399d75edce6a9f94afaa2adb577035966e）。请经常注意代码更改。
Tensorflow正在迁移到一个名为XLA的更抽象的体系结构，其范围不同。所以上面的解释将会过时。

如何可视化实际的底层计算代码？

1 个答案: