1. torch.mm()
- 功能: 用于两个2D密集矩阵的矩阵乘法(严格的矩阵乘法,不适用于向量或更高维度的张量。
- 输入: 仅支持 2D 张量,即两个矩阵的维度必须为
[m, n]
和[n, p]
,结果的维度为[m, p]
。 - 适用场景: 纯 2D 矩阵的乘法计算。
2. torch.sparse.mm()
- 功能: 只支持稀疏矩阵(sparse matrix)和密集矩阵(dense matrix)之间的矩阵乘法,输入的第一个矩阵必须是稀疏矩阵。
- 输入:
- 第一个参数是稀疏矩阵(如 COO 格式的稀疏矩阵)
- 第二个参数是密集矩阵/稀疏矩阵
- 适用场景: 当处理大规模稀疏数据(如图神经网络的邻接矩阵)时,用于减少内存开销和加速计算。
a = torch.randn(2, 3).to_sparse()
b = torch.randn(3, 2)
y = torch.sparse.mm(a, b)
'''
a:
tensor(indices=tensor([[0, 0, 0, 1, 1, 1],
[0, 1, 2, 0, 1, 2]]),
values=tensor([ 0.2511, -1.2641, 0.8681, -2.5626, -0.3517, -0.1242]),
size=(2, 3), nnz=6, layout=torch.sparse_coo)
b:
tensor([[ 1.2155, -2.2681],
[-0.0799, -1.3815],
[-1.7399, -1.3790]])
y:
tensor([[-1.1041, -0.0203],
[-2.8706, 6.4692]])
'''
3. torch.mul
- 功能: 用于张量的元素级乘法,即两个张量的相同位置上的元素相乘。按元素进行乘法运算,要求两个张量的形状相同或广播兼容。
- 输入: 可以是任意形状的张量,只要它们的形状相同或能通过广播机制匹配。
- 适用场景: 用于逐元素操作,例如标量与张量相乘,或两个形状相同的张量逐元素相乘。
a = torch.tensor([[1, 2], [3, 4]])
b = torch.tensor([[5, 6], [7, 8]])
result = torch.mul(a, b)
# result = [[1*5, 2*6], [3*7, 4*8]] = [[5, 12], [21, 32]]
4. torch.matmul
-
功能: 支持多种场景,包括:
- 1D 向量与 1D 向量相乘(内积)
- 2D 矩阵与 2D 矩阵相乘(类似
torch.mm
) - 2D 矩阵与 1D 向量相乘(矩阵和向量乘法)
- 高维张量的矩阵乘法,支持广播
-
输入: 1D、2D 或更高维度的张量,PyTorch 会根据张量的维度自动选择合适的乘法操作。
-
适用场景: 用于各种矩阵乘法场景,比
torch.mm
更通用,可以处理高维度的张量。
# Vector dot product (1D)
a = torch.tensor([1, 2])
b = torch.tensor([3, 4])
result = torch.matmul(a, b) # Dot product, result = 1*3 + 2*4 = 11
# 2D matrix multiplication (like torch.mm)
a = torch.randn(2, 3)
b = torch.randn(3, 4)
result = torch.matmul(a, b) # Shape will be (2, 4)
# 3D tensor multiplication with broadcasting
a = torch.randn(10, 3, 4)
b = torch.randn(4)
result = torch.matmul(a, b) # Shape will be (10, 3), broadcasting the 4D vector
5. torch.bmm
与torch.mm类似,但多了一个batch_size维度,矩阵张量1维度是(b×m×n),矩阵张量2维度是(b×n×p),输出维度为(b×m×p)
a = torch.randn(10, 3, 4)
b = torch.randn(10, 4, 5)
result = torch.matmul(a, b) # Shape will be (10, 3, 5)
结合批量分析 torch.matmul中
# 3D tensor multiplication with broadcasting
a = torch.randn(10, 3, 4)
b = torch.randn(4)
3D的第一个维度当作是批量,即a
有10个(3, 4)
的矩阵,b是1D,可以视为行向量或者列向量均可,在进行矩阵乘法axb时,需要满足a的列数=b的行数,广播机制后可以进行运算,size为(10,3)