¿Hay alguna forma de usar tensorflow map_fn en GPU?

Tengo un tensor A con forma [a, n] y necesito realizar una my_op op my_op con otro tensor B de forma [b, n] de modo que el tensor C resultante tenga forma [a, b].

En otras palabras: para cada subtensor en A (A [0], A 1 , … A [n]) necesito realizar un op sabio de elementos con cada subtensor en B.

Así que el tensor resultante contendría lo siguiente:

 [ [ A[0] op B[0] , A[0] op B[1], ... , A[0] op B[b] ], [ A[1] op B[0] , A[1] op B[1], ... , A[1] op B[b] ], [ ... ], [ A[a] op B[0] , A[a] op B[1], ... , A[a] op B[b] ] ] 

La única forma en que he podido encontrar que logra esto es mediante el uso nested de tf.map_fn. Así:

 import tensorflow as tf import time import numpy as np a_size = 64 b_size = 256*256 n = 256 A = tf.placeholder(tf.float32,[a_size,n]) B = tf.placeholder(tf.float32,[b_size,n]) def elementwise_op(a,b): return tf.reduce_sum(tf.multiply(a,b)) def intermediate_op(sub_a,my_b): sample_values = tf.map_fn(lambda x: elementwise_op(sub_a,x),my_b) return sample_values my_op = tf.map_fn(lambda x: intermediate_op(x,B),A) with tf.Session() as sess: a = np.random.rand(a_size,n) b = np.random.rand(b_size,n) start_time = time.time() result = sess.run (my_op,feed_dict={A:a,B:b}) print ("exec time: " ,time.time()-start_time) print (result.shape) 

El código anterior funciona bien, sin embargo, no usa muy bien la GPU (solo un 15% de utilización, según nvidia-smi ). De hecho, ¡se ejecuta un orden de magnitud más rápido cuando se usa solo la CPU ! (en mi máquina de 12 núcleos) Cuando se ejecuta utilizando la GPU, veo una utilización de GPU muy baja (~ 15%) y el 100% en uno de mis núcleos de CPU. Cuando se ejecuta solo en la CPU, veo un 100% de utilización en todos los núcleos de la CPU.

    El tiempo promedio de 5 CPU solo se ejecuta: 11.33s

    Tiempo promedio de 5 carreras de GPU: 111.88s

    La prueba anterior se ejecutó usando las imágenes oficiales de Tensorflow tensorflow/tensorflow:latest-py3 : tensorflow/tensorflow:latest-py3 (para CPU) y tensorflow/tensorflow:latest-gpu-py3 (para GPU)

    Supongo que map_fn , a través de python lambda, está obligando a que los datos se copien entre la CPU y la GPU en cada iteración, y la naturaleza anidada de la operación solo la empeora. Los comentarios en la pregunta de SO sin responder aquí sugieren que este es el caso.

    Este artículo afirma que:

    La expresión lambda es la razón principal de la baja utilización de GPU.

    Entonces, mi pregunta es: ¿hay una manera de forzar a map_fn a usar la GPU? ¿O para evitar la lambda de Python?

    Alternativamente, ¿hay alguna otra forma (quizás más de tensorflow-y) para lograr el resultado descrito anteriormente, para que la gráfica de obtención se ejecute en la GPU?

    Edición: después de ejecutar el generador de perfiles (tuve que reducir drásticamente el tamaño de los arreglos para que el generador de perfiles funcionara, porque estaba consumiendo RAM como locos), las siguientes líneas llamaron mi atención:

     node name | output bytes | total execution time | accelerator execution time | cpu execution time Mul 1.02KB (22.23%, 0.29%), 195.07ms (85.00%, 13.06%), 5.29ms (100.00%, 25.79%), 189.78ms (84.79%, 12.89%) Sum 256B (21.41%, 0.07%), 241.48ms (69.08%, 16.17%), 6.01ms (74.21%, 29.29%), 235.47ms (69.01%, 15.99%) TensorArrayScatterV3 512B (0.64%, 0.15%), 658.31ms (46.87%, 44.09%), 9.19ms (44.80%, 44.80%), 649.12ms (46.90%, 44.08%) 

    Parece que ciertas operaciones se realizan principalmente en la CPU, y solo en un hilo en eso!

    Related of "¿Hay alguna forma de usar tensorflow map_fn en GPU?"

    La construcción tf.map_fn() se puede usar con una función que ejecuta ops en GPU. De forma predeterminada, TensorFlow intentará ejecutar la mayor parte posible de la función en la GPU, y cualquier operación incompatible con GPU se ejecutará en la CPU. En su progtwig, toda la función elementwise_op() se construye a partir de operaciones compatibles con GPU, por lo que no debe haber copia adicional entre la CPU y la GPU en cada iteración.

    La causa de la baja utilización de GPU es difícil de determinar a partir de un fragmento de progtwig. Por ejemplo, si A y B son relativamente pequeños, y los está alimentando desde Python y el resultado se recupera de inmediato, es probable que la sobrecarga de copiar los datos iniciales hacia y desde la GPU domine. La mejor manera de rastrear esto es usar un generador de perfiles de GPU, que puede obtener usando tfprof o NVIDIA Visual Profiler .