python - 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() 的实现是用 C++ 编写的，它使用 Eigen(在 CPU 上)或 cuDNN 库(在 GPU 上)调用优化代码。你可以找到实现 here .

您在问题中提到的函数链(从 tf.nn.conv2d() 向下)是用于构建 TensorFlow 图的 Python 函数，但这些不调用实现。回想一下，在 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 值所需的所有操作，包括卷积本身。从这里到实现的路径有些复杂，但经过以下步骤:

1. sess.run() 调用 TensorFlow 后端来获取 conv 的值。
2. 后端修剪计算图以确定必须执行哪些节点，并将节点放置在适当的设备(CPU 或 GPU)上。
3. 指示每个设备执行其子图，使用 executor .
4. 执行器通过调用它的 Compute() 方法最终调用与卷积运算符对应的 tensorflow::OpKernel。

“Conv2D” OpKernel 已实现 here ，它的Compute()方法是here .由于此操作对于许多工作负载的性能至关重要，因此实现相当复杂，但基本思想是将计算卸载到 Eigen Tensor 库(如果在 CPU 上运行)或 cuDNN 的优化 GPU 实现。