я есть тензорA с формой [а, п], и мне нужно выполнить опmy_op с другим тензоромB формы [b, n] такой, что результирующий тензорC имеет форму [а, б].

Другими словами: длякаждый субтензор вA (A [0], A1, ... A [n]) Мне нужно выполнить поэлементную операцию скаждый субтензор вB.

Таким образом, полученный тензор будет содержать следующее:

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

Единственный способ, которым я смог найти это, достигается путем вложенного использованияtf.map_fn Таким образом:

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)

Приведенный выше код работает нормально, однако он не очень хорошо использует графический процессор (только ~ 15%, согласноnvidia-smi). На самом деле, он работает на порядок быстрее при использованиитолько процессор!! (на моем 12-ядерном компьютере) При запуске с использованием графического процессора я вижу очень низкое использование графического процессора (~ 15%) и 100% наодин моих процессорных ядер. Когда запускается только на процессоре, я вижу 100% загрузку на всех ядрах процессора.

Среднее время 5 CPU только работает:11.33s

Среднее время работы 5 GPU:111.88s

Приведенный выше тест был выполнен с использованием официальных изображений док-станции Tensorflow:tensorflow/tensorflow:latest-py3 (для процессора) иtensorflow/tensorflow:latest-gpu-py3 (для графического процессора)

Я думаю, чтоmap_fnчерез лямбду Python заставляет копировать данные назад и вперед между процессором и графическим процессором вкаждый итерация, и вложенная природа операции только ухудшает ситуацию. Комментарии в неотвеченном вопросеВот предположить, что это так.

эта статья заявляет, что:

Лямбда-выражение является основной причиной низкого использования графического процессора.

-

Итак, мой вопрос: есть ли способ заставить map_fn использовать графический процессор? Или чтобы избежать Python лямбда?

В качестве альтернативы, есть ли какой-то другой (возможно, более tenorflow-y) способ достижения результата, описанного выше, для того, чтобы получить граф для работы на GPU?

Редактировать: После запуска профилировщика (мне пришлось резко уменьшить размер массивов, чтобы заставить профилировщик работать вообще, потому что он потреблял ОЗУ как сумасшедший), следующие строки привлекли мое внимание:

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%)

Похоже, что определенные операции выполняются в основном на процессоре, и только на одном потоке!