bmm 是 batched matrix multiple 的简写,即批量矩阵乘法,矩阵是二维的,加上batch一个维度,因此该函数的输入必须是两个三维的 tensor,三个维度代表的含义分别是:(批量,行,列)。
对于 torch.bmm(tensor_a, tensor_b) 而言,
tensor_a 的 shape 为 (a, b, c)
tensor_b 的 shape 为 (d, e, f)
要求 a = d, c = e,即批量数相同,在计算时 tensor_a 的第 i 个矩阵与 tensor_b 的第 i 个矩阵作乘法,i = 1, 2, 3, ..., a。因此为了矩阵乘法能够进行,c 和 e 必须相同。计算过程如图1所示。
测试代码如下:
import torchBatchMatrix1 = torch.randn((3,4,3))
BatchMatrix2 = torch.randn((3,3,4))BatchMatrixMultiple = torch.bmm(BatchMatrix1, BatchMatrix2)print(BatchMatrixMultiple.shape)
输出为,与图1中绿色矩阵对应。