Pytorch中实现固定矩阵与批量向量逐行相乘的高效方法

实现核心依赖PyTorch原生广播机制，无需编写循环即可完成向量化运算，效率最优：

• 首先将形状为(dB, d2)的批量向量V通过unsqueeze操作增加1个中间维度，得到形状为(dB, 1, d2)的张量
• 直接将形状为(d1, d2)的固定矩阵M和扩展后的V张量执行逐元素相乘，PyTorch会自动将两个张量广播到统一形状(dB, d1, d2)，输出结果完全符合OUT[i] = M*V[i]的要求

完整代码示例 #

import torch

# 自定义维度参数
d1 = 5
d2 = 10
dB = 32

# 初始化输入张量
M = torch.randn(d1, d2)
V_batch = torch.randn(dB, d2)

# 核心实现
OUT = M * V_batch.unsqueeze(1)

# 验证输出形状
print(OUT.shape) # 输出：torch.Size([32, 5, 10])，即(dB, d1, d2)

正确性验证 #

可以通过如下代码验证运算结果完全匹配需求：

for i in range(dB):
    assert torch.allclose(OUT[i], M * V_batch[i])
print("所有批量样本结果验证通过")

方案优势 #

• 全程为原生向量化运算，无Python层循环开销，运算效率极高
• 广播逻辑由PyTorch底层实现，不会产生额外的张量拷贝，内存占用最优
• 代码简洁易读，维护成本低

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