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

Question

我有两个 3D 张量，张量 A 的尺寸为 [32,1024,128] 和张量 B 的尺寸为 [32,1024,1024] ，其中 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.tile 来代替使用 tf.expand_dims 进行平铺，ojit_code 实际上会通过多次复制原始张量来创建新的张量，而 ojit_a 在内存方面会具有更高的性能。

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