PyTorch张量广播与批量矩阵乘法维度适配方法咨询

要实现该需求，核心是利用PyTorch自动广播机制，先让(b,c)形状的数组自动扩展为(a,b,c)（自动在第0维补维度），再根据运算需求得到(a,c)结果：

• 假设输入张量为X（形状[a,b,c]），数组为Y（形状[b,c]）
• 示例代码（以对应元素相乘后沿b维度求和为例）：

  # 广播后逐元素相乘，再沿b维度求和得到(a,c)输出
  result = (X * Y).sum(dim=1)
  

针对你的具体场景，有两种简洁实现方式：

方式1：调整B的形状后使用torch.matmul #

将B的形状从(416,20)扩展为(416,20,1)，torch.matmul会自动处理batch维度的矩阵乘法：

import torch

A = torch.randn(416, 20, 3)
B = torch.randn(416, 20)

# 扩展维度后执行矩阵乘法，再移除最后一个冗余维度
C = torch.matmul(A, B.unsqueeze(-1)).squeeze(-1)
print(C.shape)  # 输出 torch.Size([416, 3])

方式2：使用torch.einsum（维度映射更直观） #

爱因斯坦求和符号可直接指定维度运算关系，可读性更强：

C = torch.einsum('abc,ab->ac', A, B)
print(C.shape)  # 输出 torch.Size([416, 3])

这里'abc,ab->ac'表示：对每个batch的a维度，将A的b维度与B的b维度相乘求和，最终保留a和c维度，完全匹配你的需求。

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