pytorch矩阵乘法总结

1. element-wise（*） 按元素相乘，支持广播，等价于torch.mul()

a = torch.tensor([[1, 2], 
								 [3, 4]])
b = torch.tensor([[2, 3], 
								 [4, 5]])
c = a*b  # 等价于torch.mul(a,b)
# tensor([[ 2,  6],
#        [12, 20]])
a * torch.tensor([1, 2]) # 广播,等价于torch.mul(a, torch.tensor([1, 2]))
# tensor([[1, 4],
#        [3, 8]])
a * torch.tensor([2])  # 乘以单个元素，依然是广播,等价于torch.mul(a, torch.tensor([2]))
# tensor([[1, 4],
#         [3, 8]])

2. torch.mm(), 实现矩阵叉乘，即第一个矩阵的列数等于第二个矩阵的行数，不支持广播。

a = torch.arange(9).reshape(3, 3)
b = torch.arange(12).reshape(3, 4)
torch.mm(a, b)
#tensor([[ 20,  23,  26,  29],
#        [ 56,  68,  80,  92],
#        [ 92, 113, 134, 155]])

3. torch.bmm(),三维矩阵乘法，第一个维度是batch_size的大小。两个矩阵的batch_size大小要相等。

a = torch.arange(9).reshape(1,3, 3)
b = torch.arange(12).reshape(1,3, 4)
torch.bmm(a, b)
# tensor([[[ 20,  23,  26,  29],
#         [ 56,  68,  80,  92],
#         [ 92, 113, 134, 155]]])

4. torch.matmul(), 混合矩阵乘法

4.1 都是一维，进行点积运算

4.2 都是二维，返回矩阵叉乘

4.3 向量乘以矩阵

4.4 矩阵乘以向量

4.5 一维向量乘以多维矩阵

# 如果两个参数都是至少1个维度，并且至少一个参数是3维，那么返回一个batched矩阵乘法。
# 一维向量乘以多维矩阵
c = torch.arange(6).reshape(1,2, 3)
d = torch.tensor([2, 1])

e = torch.matmul(d, c)  # d前置增加一个维度，然后在batch_size维度进行广播
e
# tensor([[3, 6, 9]])

4.6 多维矩阵乘以一维向量

e = torch.arange(2).reshape(1, 1, 2)
# 多维维矩阵乘以向量，向量在后面增加一个维度，然后在batch_size维度进行广播
f = torch.tensor([2, 4])  # 变成[[2], [4]]
g = torch.matmul(e, f)
g
# tensor([[4]])

4.7 最后两个维度没有一维向量

最后两个维度进行矩阵叉乘，前面的维度进行广播进行计算。
感兴趣的可以看看文档：https://pytorch.org/docs/stable/generated/torch.matmul.html#torch.matmul

5 torch.einsum()

参考该链接：https://zhuanlan.zhihu.com/p/434232512

6 Reference:

[1] https://huaweicloud.csdn.net/63807229dacf622b8df8865b.html
[2] https://pytorch.org/docs/stable/generated/torch.matmul.html#torch.matmul