Tensorflow：tf.nn.conv2d 实际执行在哪里？

Question

我很好奇 tf.nn.conv2d(...) 的 Tensorflow 实现。要调用它，只需运行 tf.nn.conv2d(...)。但是，我正在尝试查看它的执行位置。代码如下（其中箭头表示最终调用的函数）：

tf.nn.conv2d(...) -> tf.nn_ops.conv2d(...) -> tf.gen_nn_ops.conv2d(...) -> _op_def_lib.apply_op("Conv2D", ...) -> ?

我熟悉 Tensorflow 的 LSTM 实现以及根据自己的需要轻松操作它们的能力。执行conv2d()计算的函数是否写在Python中，如果有，它在哪里？我可以看到步幅在哪里以及如何执行吗？

Answer 1

TL;DR:执行tf.nn.conv2d() is written in C++, which invokes optimized code using either Eigen (on CPU) or the cuDNN library (on GPU). You can find the implementation here.

您在问题中提到的函数链（从 tf.nn.conv2d() 往下）是 Python 用于构建 TensorFlow 图的函数，但是这些不调用实现。回想一下，在 TensorFlow 中，您首先 build a symbolic graph, then execute it.

tf.nn.conv2d() 的实现仅在您调用 Session.run() 并传递一个 Tensor 时才会执行，其值取决于某些卷积的结果。例如：

input = tf.placeholder(tf.float32)
filter = tf.Variable(tf.truncated_normal([5, 5, 3, 32], stddev=0.1)
conv = tf.nn.conv2d(input, filter, strides=[1, 1, 1, 1], padding='SAME')

result = sess.run(conv, feed_dict={input: ...})  # <== Execution happens here.

调用 sess.run(...) 告诉 TensorFlow 运行计算 conv 的值所需的所有操作，包括卷积本身。从这里到实现的路径有点复杂，但经过以下步骤：

sess.run() 调用 TensorFlow 后端来获取 conv.
后端 p运行计算图计算出必须执行的节点，并将节点放置在适当的设备（CPU 或 GPU）上。
指示每个设备执行其子图，使用 executor。
执行器最终通过调用其Compute()方法调用对应于卷积运算符的tensorflow::OpKernel。

"Conv2D" OpKernel 已实现 here, and its Compute() method is here。因为此操作对于许多工作负载来说都是性能关键，所以实现起来相当复杂，但基本思想是将计算卸载到 Eigen Tensor 库（如果运行ning on CPU）或 cuDNN 的优化的 GPU 实现。

Answer 2

TensorFlow 程序由两个独立的部分组成：

构建计算图。

tf.nn.conv2d(...) -> tf.nn_ops.conv2d(...) -> tf.gen_nn_ops.conv2d(...) -> _op_def_lib.apply_op("Conv2D", ...) -> graph.create_op -> 将 op 注册到图

运行计算图。

sess = tf.Session(target) -> sess.run(conv2d) -> master p运行e full graph to client graph -> master split client graph by task to graph partition -> 将图分区注册到 worker -> worker 按设备拆分子图到图分区 -> 然后 master 通知所有 worker 到运行图分区 -> worker 通知所有设备到运行图分区 -> 执行器将运行在设备上按拓扑排序操作。

对于op之一，执行器将调用内核实现来计算op。

tf.nn.conv2d() 的内核实现是用 C++ 编写的，它使用 Eigen（在 CPU 上）或 cuDNN 库（在 GPU 上）调用优化代码。

Tensorflow：tf.nn.conv2d 实际执行在哪里？

Tensorflow: Where is tf.nn.conv2d Actually Executed?

python

machine-learning

tensorflow