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

Question

我有两个 3D 张量，维度为 [32,1024,128] 的张量 A 和维度为 [32,1024,1024] 的张量 B ，其中 32 是批量大小。对于某个样本，我想将矩阵 A 的每一行与矩阵 B 的每个元素相乘。这样输出张量维度将为 [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.expand_dims，而不是使用 tf.tile 进行平铺，后者实际上是通过多次复制原始张量来创建新张量，它在内存方面的性能更高。

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