python - seq2seq - Tensorflow: ¿Dónde se ejecuta realmente tf.nn.conv2d?

tensorflow seq2seq

Question

Tengo curiosidad sobre la implementación tf.nn.conv2d(...) de tf.nn.conv2d(...) . Para llamarlo, uno simplemente ejecuta tf.nn.conv2d(...) . Sin embargo, voy por el agujero del conejo tratando de ver dónde se ejecuta. El código es el siguiente (donde la flecha indica la función que finalmente llama):

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

Estoy familiarizado con la implementación de LSTM de Tensorflow y la capacidad de manipularlos fácilmente según lo considere necesario. Es la función que realiza el cálculo conv2d() escrito en Python, y si es así, ¿dónde está? ¿Puedo ver dónde y cómo se ejecutan los pasos?

Answer 1

Los programas TensorFlow consisten en dos secciones discretas:

• Construyendo el grafico computacional.

tf.nn.conv2d (...) -> tf.nn_ops.conv2d (...) -> tf.gen_nn_ops.conv2d (...) -> _op_def_lib.apply_op ("Conv2D", ...) -> graph.create_op -> register op into graph

• Ejecutando el grafo computacional.

sess = tf.Session (target) -> sess.run (conv2d) -> master podar el gráfico completo al gráfico del cliente -> master dividir el gráfico del cliente por tarea para graficar la partición -> registrar la partición del gráfico al trabajador -> trabajador dividir el subgrafo por dispositivo para graficar partición -> luego maestro notificar a todos los trabajadores para ejecutar particiones gráficas -> trabajador notificar a todos los dispositivos para ejecutar particiones gráficas -> ejecutor ejecutará operaciones por orden topológico en el dispositivo.

Para uno de op, el ejecutor invocará kernel implement para calcular para el op.

El implemento kernel de tf.nn.conv2d () está escrito en C ++, que invoca código optimizado usando Eigen (en la CPU) o la biblioteca cuDNN (en la GPU).

Answer 2

TL; DR: la implementación de tf.nn.conv2d() está escrita en C ++, que invoca el código optimizado usando Eigen (en la CPU) o la biblioteca cuDNN (en la GPU). Puede encontrar la implementación aquí .

La cadena de funciones que mencionas en la pregunta (desde tf.nn.conv2d() hacia abajo) son funciones de Python para construir un gráfico de TensorFlow, pero estas no invocan la implementación. Recuerde que, en TensorFlow, primero crea un gráfico simbólico y luego lo ejecuta .

La implementación de tf.nn.conv2d() solo se ejecuta cuando se llama a Session.run() pasando un Tensor cuyo valor depende del resultado de alguna convolución. Por ejemplo:

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.

La invocación de sess.run(...) le dice a TensorFlow que ejecute todas las operaciones que se necesitan para calcular el valor de conv , incluida la propia convolución. El camino desde aquí hasta la implementación es algo complicado, pero sigue los siguientes pasos:

1. sess.run() llama al backend de TensorFlow para recuperar el valor de conv .
2. El backend corta el gráfico de cálculo para determinar qué nodos deben ejecutarse y coloca los nodos en los dispositivos apropiados (CPU o GPU).
3. Cada dispositivo tiene instrucciones de ejecutar su subgrafo, usando un ejecutor .
4. El ejecutor finalmente invoca el tensorflow::OpKernel que corresponde al operador de convolución, llamando a su método Compute() .

El "Conv2D" se implementa aquí , y su método Compute() está aquí . Como esta operación es crítica para muchas cargas de trabajo, la implementación es bastante complicada, pero la idea básica es que el cálculo se descargue a la biblioteca Eigen Tensor (si se ejecuta en la CPU) o a la implementación de la GPU optimizada de cuDNN.