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

Question

我对Tensorflow的实现感到好奇tf.nn.conv2d(...).要调用它,只需运行即可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上)调用优化代码.你可以在这里找到实现.

您在问题中提到的函数链(从tf.nn.conv2d()下面开始)是用于构建 TensorFlow图的Python函数,但这些函数不会调用实现.回想一下,在TensorFlow中,首先构建一个符号图,然后执行它.

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. 每个设备被指示执行其子,用执行.
4. 执行程序最终tensorflow::OpKernel通过调用其Compute()方法来调用与卷积运算符相对应的.

该"Conv2D"OpKernel实现这里,其Compute()方法是在这里.由于此操作对于许多工作负载而言性能至关重要,因此实现非常复杂,但基本思想是将计算卸载到Eigen Tensor库(如果在CPU上运行)或cuDNN的优化GPU实现.