我有两个3D张量，尺寸为[32,1024,128]的张量A和尺寸为[32,1024,1024]的张量B，其中32是批处理大小。对于某个样本，我想将矩阵A的每一行与矩阵B的每个元素相乘。这样输出张量维数将为[32,1024,1024,128]。我尝试使用tf.tile将两个张量都转换为4D并使用逐元素乘法。但这给了我内存不足的错误。我尝试将批次大小减小为4，但存在相同的问题。对此的任何帮助将不胜感激。
答案 0 :(得分：3)
您的问题确实需要很大一部分内存。这是一个演示，其中我使用了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）数组。 / p>
您可以使用tf.expand_dims来代替传统的张量，而使用tf.tile进行分块实际上会创建一个新的张量，而PMD的性能在内存方面会更高。
这些是我首先要研究的两行优化。