torch 高维矩阵乘法分析，一文说透

简介

一提到矩阵乘法，大家对于二维矩阵乘法都很了解，即 A 矩阵的行乘以 B 矩阵的列。
但对于高维矩阵乘法可能就不太清楚，不知道高维矩阵乘法是怎么在计算。

建议使用torch.matmul 做矩阵乘法，其支持向量乘法 和 二维、乃至多维的矩阵乘法。

向量乘法

a1 = torch.tensor([1, 2])
res1 = torch.matmul(a1, a1)
print(res1)
print(res1.shape)

输出：

tensor(5)
torch.Size([])

torch 也支持使用 @ 完成乘法操作

二维矩阵乘法

a2 = torch.tensor([[1, 2]])
res2 = torch.matmul(a2, a2.transpose(-2, -1))
print(res2)
print(res2.shape)

输出：

tensor([[5]])
torch.Size([1, 1])

torch.mm 与 @ 也可以做二维矩阵乘法：

• a2 @ a2.transpose(-2, -1)
• torch.mm(a2, a2.transpose(-2, -1))

三维矩阵乘法

torch.bmm 支持三维矩阵乘法，不支持更高维度的矩阵乘法

a3 = torch.randn(2, 3, 2)
res3 = torch.bmm(a3,a3.transpose(-1, -2)
)
print(res3)
print(res3.shape)

输出：

tensor([[[ 4.5979,  0.6648,  2.9231],[ 0.6648,  0.1155,  0.4713],[ 2.9231,  0.4713,  1.9805]],[[ 1.0323,  1.8212, -0.3546],[ 1.8212,  3.5445, -0.3834],[-0.3546, -0.3834,  0.2988]]])
torch.Size([2, 3, 3])

a3 的 shape是(2, 3, 2)，a3 底层的两个维度做转置之后变成(2, 2, 3)，才可以做矩阵乘法。
可以发现第一位的数字都是2。高维矩阵做乘法的时候，除了最后两个维度，高维矩阵前面的维度两个矩阵要保持一致。

torch.randn(2, 3, 2) @ torch.randn(3, 2, 3)

虽然上述两个矩阵，在最后两个维度满足矩阵运算的条件，但是第一个维度两个矩阵的值不一样，所以不能做矩阵乘法。

广播

但是发现：

t1 = torch.randn(1, 3, 2)
t2 = torch.randn(3, 2, 3)
t1 @ t2

输出：

tensor([[[-0.6557,  1.0518,  0.3055],[-0.2876, -2.5104, -1.4417],[ 1.4447, -0.1799,  0.4602]],[[ 0.2971,  0.0060, -0.2612],[-0.9089,  1.0824,  0.7131],[ 0.0929, -0.7898, -0.0199]],[[ 0.0027,  1.2031,  0.1543],[-0.5603, -1.8567, -0.1302],[ 0.3978, -0.9356, -0.1977]]])

理论上两个矩阵的高维度的shape不一样，就不可以做矩阵乘法。但上述 t1与 t2可以做矩阵乘法。这是因为 t1 的第一个维度是1，就会自动做广播。

广播的效果类似于，把 t1 在第一个维度复制成与t2一样，第一个维度都变成3。
在下述使用 concat完成复制工作，再做矩阵乘法，发现可以得到上述一样的结果。

torch.concat((t1, t1, t1)) @ t2

输出：

tensor([[[-0.6557,  1.0518,  0.3055],[-0.2876, -2.5104, -1.4417],[ 1.4447, -0.1799,  0.4602]],[[ 0.2971,  0.0060, -0.2612],[-0.9089,  1.0824,  0.7131],[ 0.0929, -0.7898, -0.0199]],[[ 0.0027,  1.2031,  0.1543],[-0.5603, -1.8567, -0.1302],[ 0.3978, -0.9356, -0.1977]]])

高维矩阵乘法

矩阵乘法只会在最后两个维度，用A矩阵的行乘以B矩阵的列。
其他的维度都是对应位置的数据，互相做乘法（类似向量乘法）。

high_matrix1 = torch.randn(2, 3, 4, 5)
high_matrix2 = torch.randn(2, 3, 5, 4)
high_result = high_matrix1 @ high_matrix2

把最后两个维度看成一个点。更高的维度的矩阵乘法，可想象为两个矩阵对应位置的点相乘。

比如，shape(2, 3, 4, 5)与shape(2, 3, 5, 4)的矩阵相乘，若把最后两个维度看成一个点。就可以类比为 (2, 3) 与 (2, 3)的两个矩阵做向量乘法，就是对应位置的点做乘法。

如下面的运行结果所示。针对两个矩阵，在高维空间中，选取（1，2）对应的小矩阵数据做矩阵乘法得到的结果。与两个矩阵乘法的结果对应（1，2）的值是一样的。

(high_matrix1[1][2] @  high_matrix2[1][2]) == high_result[1][2]

输出：

tensor([[True, True, True, True],[True, True, True, True],[True, True, True, True],[True, True, True, True]])

开源

https://github.com/JieShenAI/csdn/blob/main/25/06/torch_matmul/run.ipynb