如何在PyTorch中实现4D矩阵与2D矩阵的批量乘法？

要实现尺寸为(A,B,C,D)的4D张量和(D,C)的2D矩阵的批量乘法，最终得到(A,B,C,C)的结果，有两种高效的实现方式：

方法一：直接使用torch.matmul()（推荐） #

PyTorch的matmul支持自动广播批量维度，无需手动调整张量形状，代码简洁且性能最优。

代码示例 #

import torch

# 定义维度参数
A, B, C, D = 2, 3, 4, 5

# 创建输入张量
tensor_4d = torch.randn(A, B, C, D)  # 4D张量：(A,B,C,D)
matrix_2d = torch.randn(D, C)        # 2D矩阵：(D,C)

# 执行批量矩阵乘法
result = torch.matmul(tensor_4d, matrix_2d)

# 验证结果形状
print(result.shape)  # 输出: torch.Size([2, 3, 4, 4]) → 对应(A,B,C,C)

原理说明 #

matmul会将高维张量的前N-2个维度视为批量维度，仅对最后两个维度执行矩阵乘法。这里4D张量的批量维度是(A,B)，每个批量元素是(C,D)的子矩阵；2D矩阵会被自动广播为匹配批量维度的(A,B,D,C)张量，最终每个(C,D)子矩阵与(D,C)矩阵相乘得到(C,C)结果，堆叠后形成(A,B,C,C)的输出。

方法二：结合torch.bmm()与形状调整（适合理解逻辑） #

如果需要更直观地控制批量操作流程，可以将4D张量reshape为3D，配合bmm（批量矩阵乘法）实现：

代码示例 #

import torch

A, B, C, D = 2, 3, 4, 5
tensor_4d = torch.randn(A, B, C, D)
matrix_2d = torch.randn(D, C)

# 将4D张量转为3D批量张量：(A*B, C, D)
tensor_3d = tensor_4d.reshape(-1, C, D)
# 将2D矩阵扩展为3D并匹配批量数量：(A*B, D, C)
matrix_3d = matrix_2d.unsqueeze(0).repeat(A*B, 1, 1)

# 执行批量矩阵乘法
result_3d = torch.bmm(tensor_3d, matrix_3d)
# 将结果转回4D：(A,B,C,C)
result = result_3d.reshape(A, B, C, C)

print(result.shape)  # 输出: torch.Size([2, 3, 4, 4])

原理说明 #

bmm要求输入为两个3D张量，形状分别为(N, M, K)和(N, K, L)，输出为(N, M, L)。这里我们把原4D张量的前两个维度合并为批量维度N=A*B，再将2D矩阵重复N次匹配批量大小，最后将结果拆分回原有的(A,B)维度即可。

内容的提问来源于stack exchange，提问作者장민규