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的實現以及在人們認為合適的情況下輕松操作它們的能力。 是用於執行用Python編寫的conv2d()計算的函數，如果是，它在哪里？ 我可以看到步幅的執行位置和方式嗎？

Answer 1

TL; DR： tf.nn.conv2d()的實現是用C ++編寫的，它使用Eigen（在CPU上）或cuDNN庫（在GPU上）調用優化代碼。 你可以在這里找到實現。

您在問題中提到的函數鏈（來自tf.nn.conv2d() ）是用於構建 TensorFlow圖的Python函數，但這些函數不會調用實現。 回想一下，在TensorFlow中，首先構建一個符號圖，然后執行它 。

tf.nn.conv2d()的實現僅在調用Session.run()傳遞Tensor時執行， 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. 執行程序最終通過調用其Compute()方法調用與卷積運算符對應的tensorflow::OpKernel 。

這里實現了 "Conv2D" OpKernel，它的Compute()方法就在這里 。 由於此操作對於許多工作負載而言性能至關重要，因此實現非常復雜，但基本思想是將計算卸載到Eigen Tensor庫（如果在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 prune full graph to client graph - > master split client graph by task to graph partition - > register graph partition to worker - > worker split subgraph by device to graph partition - > then master通知所有worker運行圖分區 - > worker通知所有設備運行圖分區 - > executor將通過設備上的拓撲排序運行ops。

對於其中一個，執行程序將調用內核工具來為op進行計算。

tf.nn.conv2d（）的內核實現是用C ++編寫的，它使用Eigen（在CPU上）或cuDNN庫（在GPU上）調用優化代碼。