Как умножить каждый ряд трехмерной матрицы на каждый элемент другой трехмерной матрицы в тензорном потоке? - PullRequest
Новая
       15

Как умножить каждый ряд трехмерной матрицы на каждый элемент другой трехмерной матрицы в тензорном потоке?

2 голосов
/ 02 мая 2019

У меня есть два 3D-тензора, тензор A с размером [32,1024,128] и тензор B с размером [32,1024,1024], где 32 - размер партии.Для определенной выборки я хочу умножить каждую строку матрицы A на каждый элемент матрицы B. Таким образом, размер выходного тензора будет [32,1024,1024,128].Я попытался преобразовать оба тензора в 4D с помощью tf.tile и использовал поэлементное умножение.но это дает мне ошибку памяти.Я попытался уменьшить размер пакета до 4, но та же проблема.Будем благодарны за любую помощь в этом.

1 Ответ

3 голосов
/ 02 мая 2019

Вашей проблеме действительно требуется значительная часть памяти. Вот демонстрация, где я использовал 2 примера вместо всех примеров в batch_size = 32, 

# input arrays to work with
In [2]: A = np.random.random_sample([32,1024,128])
In [3]: B = np.random.random_sample([32,1024,1024])

# inspect their memory usage

In [12]: A.nbytes/1000000
Out[12]: 33.554432   # ~ 33.5 Mb

In [13]: B.nbytes/1000000
Out[13]: 268.435456  # ~ 268 Mb

# your desired multiplication
In [14]: res = B[:2, ..., np.newaxis] * A[:2, :, np.newaxis, ...]

# desired shape of the output
In [15]: res.shape
Out[15]: (2, 1024, 1024, 128)

# inspect memory usage
In [16]: res.nbytes/1000000
Out[16]: 2147.483648  # ~ 2.1 GB

Я использовал float64 для этих массивов. Если вы не можете позволить себе такое требование к памяти, то одной из идей уменьшить использование памяти, избегая при этом ошибок Out Of Memory, было бы уменьшить ваши массивы и использовать массивы с одинарной точностью (т.е. float32).

Вместо разбиения на листы с использованием tf.tile, который фактически создает новый тензор путем многократного копирования оригинального тензора, вы можете использовать tf.expand_dims, который был бы более производительным с точки зрения памяти.

Это две линии оптимизации, которые я рассмотрю в первую очередь.

...