PyTorch不同维度张量相乘：如何优雅实现a每行乘b对应标量

PyTorch内置的广播机制可以直接满足你的需求，不需要写循环逐行计算，性能和可读性都更好，下面是几种常用的实现方式：

• 方法1：用索引语法直接给b新增维度（写法最简洁）

c = a * b[:, None]

这里b[:, None]会把形状为(40,)的b调整为(40, 1)，触发广播机制后会自动扩展为和a相同的(40, 6)形状，每一行的6个元素都对应b原来的标量，相乘后正好是你要的行维度对应相乘的效果。

• 方法2：用unsqueeze方法显式升维（可读性更强）

c = a * b.unsqueeze(dim=1)

和方法1效果完全等价，unsqueeze(1)的作用就是在第1维（0为起始索引）新增一个长度为1的维度。

• 方法3：用爱因斯坦求和接口实现（逻辑更直观）
  如果你熟悉einsum语法，可以直接指定维度匹配规则，不需要考虑广播逻辑：

c = torch.einsum('ij,i->ij', a, b)

你可以用下面的代码验证结果和你逐行计算的结果完全一致：

c_loop = torch.stack([a[i] * b[i] for i in range(a.shape[0])])
print(torch.allclose(c, c_loop)) # 输出为True代表结果一致

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