¿Hay alguna manera de usar tensorflow map_fn en la GPU?

Tengo un tensorA con forma [a, n] y necesito realizar una operaciónmy_op con otro tensorB de forma [b, n] tal que el tensor resultanteC tiene forma [a, b].

En otras palabras: paracada subtensor enA (A [0], A1, ... A [n]) Necesito realizar un elemento sabio concada subtensor enB.

Entonces 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 que he podido encontrar que logra esto es a través del uso anidado detf.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 la GPU muy bien (solo ~ 15% de utilización, segúnnvidia-smi) De hecho, se ejecuta un orden de magnitud más rápido cuando se usasolamente la CPU!! (en mi máquina de 12 núcleos) Cuando se ejecuta con la GPU, veo una utilización de GPU muy baja (~ 15%) y 100% enuno de mis núcleos de CPU. Cuando se ejecuta solo en la CPU, veo una utilización del 100% en todos los núcleos de la CPU.

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

Tiempo promedio de 5 ejecuciones de GPU:111.88s

La prueba anterior se ejecutó utilizando las imágenes oficiales del acoplador Tensorflow:tensorflow/tensorflow:latest-py3 (para CPU) ytensorflow/tensorflow:latest-gpu-py3 (para GPU)

Mi conjetura es quemap_fn, a través de python lambda, está obligando a los datos a copiarse entre la CPU y la GPU encada iteración, y la naturaleza anidada de la operación solo lo empeora. Los comentarios en la pregunta SO sin respuestaaquí Sugerir 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 alguna forma de obligar a map_fn a usar la GPU? ¿O para evitar la lambda Python?

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

Editar: Después de ejecutar el generador de perfiles (tuve que reducir drásticamente el tamaño de las matrices para que el generador de perfiles se ejecutara, porque estaba consumiendo RAM como loco), las siguientes líneas me llamaron la 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!

Respuestas a la pregunta(1)

Su respuesta a la pregunta