如何在tensorflow上将3D矩阵的每一行与另一个3D矩阵的每个元素相乘？

Question

我有两个三维张量，张量A的维数为[32,1024,128]，张量B的维数为[32,1024,1024]，其中32是批大小。对于某个样本，我想用矩阵B的每个元素乘以矩阵A的每一行，这样输出的张量维数将是[32,1024,1024,128]。我尝试使用tf.tile将两个张量转换为4D，并使用元素乘法。但这让我的记忆错误消失了。我试着将批处理大小缩小到4，但问题是相同的。如对此有任何帮助，将不胜感激。

Answer 1

您的问题确实需要很大一部分内存。下面是一个演示，我使用了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。

这是我首先要研究的两行优化。