如何实现PyTorch对应张量的矩阵-列向量乘法？需保留梯度

首先明确你的张量形状：

• 第一个张量（记为matrices）：shape是(5, 2, 2)，代表5个2x2矩阵
• 第二个张量（记为vectors）：shape是(5, 2)，代表5个2维列向量（PyTorch中默认以行向量形式存储，需调整维度适配矩阵乘法）

要实现对应位置的矩阵-向量乘法且保留梯度，你可以用以下两种简洁的方法：

方法1：使用torch.bmm（批量矩阵乘法） #

torch.bmm专门处理批量矩阵乘法，要求输入为3维张量（(batch, n, m) 和 (batch, m, p)），输出为(batch, n, p)。我们只需把向量张量扩展为3维（增加最后一个维度转为列向量），运算后再压缩多余维度即可：

import torch

# 初始化带梯度的张量
matrices = torch.tensor([[[8.1776, 0.6560],
                          [0.6560, 2.3653]],

                         [[8.1776, 0.6560],
                          [0.6560, 2.3104]],

                         [[8.9871, 0.6560],
                          [0.6560, 2.2535]],

                         [[1.3231, 0.6560],
                          [0.6560, 2.3331]],

                         [[4.8677, 0.6560],
                          [0.6560, 2.2935]]], requires_grad=True)

vectors = torch.tensor([[-0.1836, -0.9153],
                        [-0.1836, -0.8057],
                        [-0.2288, -0.6442],
                        [ 0.1017, -0.8555],
                        [-0.0175, -0.7637]], requires_grad=True)

# 将向量转为(5,2,1)的列向量批量格式
vectors_col = vectors.unsqueeze(-1)
# 执行批量矩阵乘法，得到(5,2,1)的结果
result = torch.bmm(matrices, vectors_col)
# 压缩最后一个冗余维度，得到(5,2)的目标张量
result = result.squeeze(-1)

# 验证梯度保留：执行反向传播测试
loss = result.sum()
loss.backward()
print(matrices.grad is not None)  # 输出True，说明梯度正常保留
print(vectors.grad is not None)   # 输出True，说明梯度正常保留

方法2：使用@运算符配合维度调整 #

@运算符支持广播机制，只要维度匹配即可完成运算，同样需要先将向量转为列向量维度：

# 调整向量维度后执行矩阵乘法
result = matrices @ vectors.unsqueeze(-1)
result = result.squeeze(-1)

# 梯度保留验证逻辑与方法1一致
loss = result.sum()
loss.backward()

为什么之前的方法无效？ #

• *是元素级乘法，并非矩阵乘法，完全不符合你的需求；
• 直接使用matrices @ vectors会因维度不匹配报错：matrices为(5,2,2)，vectors为(5,2)，PyTorch无法自动对齐批量维度下的矩阵-向量乘法规则；
• torch.split是张量拆分工具，完全没必要用在这个场景，只会增加代码复杂度。

内容的提问来源于stack exchange，提问作者thedumbkid