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

【问题标题】：How to multiply each row of a 3D matrix with each element of another 3D matrix on tensorflow?如何将 3D 矩阵的每一行与 tensorflow 上另一个 3D 矩阵的每个元素相乘？
【发布时间】：2019-09-20 12:41:05
【问题描述】：

我有两个 3D 张量，尺寸为 [32,1024,128] 的张量 A 和尺寸为 [32,1024,1024] 的张量 B，其中 32 是批量大小。对于某个样本，我想将矩阵 A 的每一行与矩阵 B 的每个元素相乘。这样输出张量维度将为[32,1024,1024,128]。我尝试使用 tf.tile 将两个张量转换为 4D 并使用元素乘法。但它给了我内存不足的错误。我尝试将批量大小减少到 4，但同样的问题。对此的任何帮助将不胜感激。

【问题讨论】：

标签： python tensorflow machine-learning neural-network deep-learning

【解决方案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.expand_dims，而不是使用tf.tile（实际上通过多次复制原始张量来创建新张量）进行平铺，这在内存方面会更高效。

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

【讨论】：