Tensorflow: Где выполняется tf.nn.conv2d Выполнено?

Мне интересна реализация 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", ...) -> ?

Я знаком с реализацией LSTM от Tensorflow и возможностью легко манипулировать ими, как можно сочтет нужным. Является ли функция, которая выполняет вычисление conv2d(), написанное на Python, и если да, то где это? Могу ли я увидеть, где и как выполняются шаги?

ОТВЕТЫ

Ответ 1

TL; DR: Реализация tf.nn.conv2d() написана на С++, которая вызывает оптимизированный код, используя либо Eigen (on CPU) или библиотеки cuDNN (на GPU). Вы можете найти реализацию здесь.

Цепочки функций, которые вы упомянули в вопросе (от tf.nn.conv2d() вниз), являются функциями Python для построения графика TensorFlow, но они не вызывают реализацию. Напомним, что в 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 о запуске всех ops, которые были выбраны для вычисления значения conv, включая саму свертку. Путь отсюда к реализации несколько сложный, но проходит следующие шаги:

  • sess.run() вызывает внутренний ресурс TensorFlow для получения значения conv.
  • Бэкэнд выровняет график вычислений, чтобы определить, какие узлы должны быть выполнены, и размещает узлы на соответствующих устройствах (CPU или GPU).
  • Каждому устройству предлагается выполнить его подграф, используя executor.
  • Исполнитель в конечном итоге вызывает tensorflow::OpKernel, который соответствует оператору свертки, вызывая его метод Compute().

"Conv2D" OpKernel реализован здесь, а метод Compute() здесь. Поскольку эта операционная система является критически важной для многих рабочих нагрузок, реализация довольно сложная, но основная идея заключается в том, что вычисление выгружается либо в библиотеку Eigen Tensor (если выполняется на процессоре), либо в оптимизированную реализацию GPU cuDNN.

Ответ 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) → мастер обрезать полный график для графиков клиентов → мастер-график разделения клиентов по задаче на графический раздел → зарегистрировать граф-раздел для рабочего → раскол работника subgraph by device to graph partition → затем мастер уведомляет всех работников о запуске разделов графа → работник уведомляет все устройства о запуске разделов графа → исполнитель запускает ops топологической сортировкой на устройстве.

Для одного из op исполнитель вызовет реализацию ядра для вычисления op.

Реализация ядра tf.nn.conv2d() написана на С++, которая вызывает оптимизированный код с использованием либо Eigen (на CPU), либо библиотеки cuDNN (на графическом процессоре).