torch.bmm功能解读

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所示。
图1. bmm计算过程

测试代码如下：

import torch

BatchMatrix1 = torch.randn((3,4,3))
BatchMatrix2 = torch.randn((3,3,4))

BatchMatrixMultiple = torch.bmm(BatchMatrix1, BatchMatrix2)

print(BatchMatrixMultiple.shape)

输出为，与图1中绿色矩阵对应。