Eu tenho um tensorA com a forma [a, n] e preciso executar uma operaçãomy_op com outro tensorB de forma [b, n] tal que o tensor resultanteC tem forma [a, b].

Em outras palavras: Paracada subtensor emA (A [0], A1, ... A [n]) Eu preciso executar um elemento op sábio comcada subtensor emB.

Portanto, o tensor resultante deve conter o seguinte:

[ [ 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] ] ]

A única maneira de descobrir que consegue isso é através do uso aninhado detf.map_fn Portanto:

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)

O código acima funciona bem, no entanto, ele não usa a GPU muito bem (apenas ~ 15% de utilização, de acordo comnvidia-smi) De fato, ele executa uma ordem de magnitude mais rapidamente ao usarsó a CPU!! (na minha máquina de 12 núcleos) Ao executar a GPU, vejo uma utilização muito baixa da GPU (~ 15%) e 100% na1 dos meus núcleos de CPU. Quando executado apenas na CPU, vejo 100% de utilização em todos os núcleos da CPU.

O tempo médio de 5 CPU é executado apenas:11.33s

Tempo médio de 5 execuções de GPU:111.88s

O teste acima foi executado usando as imagens oficiais da janela de encaixe Tensorflow:tensorflow/tensorflow:latest-py3 (para CPU) etensorflow/tensorflow:latest-gpu-py3 (para GPU)

Meu palpite é quemap_fn, via python lambda, está forçando os dados a serem copiados para frente e para trás entre a CPU e a GPU emcada iteração e a natureza aninhada da operação só piora a situação. Os comentários na pergunta SO sem respostaaqui sugira que seja esse o caso.

Este artigo afirma que:

A expressão lambda é o principal motivo da baixa utilização da GPU.

-

Então, minha pergunta é: Existe uma maneira de forçar o map_fn a usar a GPU? Ou para evitar o lambda Python?

Como alternativa, existe alguma outra maneira (talvez mais tensorflow-y) de alcançar o resultado descrito acima, para que o gráfico get seja executado na GPU?

Editar: Depois de executar o criador de perfil (eu tive que reduzir drasticamente o tamanho das matrizes para que o criador funcionasse, porque estava consumindo RAM como um louco), as seguintes linhas chamaram minha atenção:

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 certas operações estão sendo feitas principalmente na CPU, e apenas em um thread!