如何将PyTorch张量a的最后维度元素与张量b对应元素相乘？

你需要利用PyTorch的广播机制来实现这个需求。核心是调整张量b的维度，让它能和a的维度对齐，从而逐位置对a最后一维的所有元素进行乘法操作。

实现方法 #

因为a的形状是(2,3,5)，b的形状是(2,3)，我们需要给b增加一个长度为1的最后维度，使其形状变为(2,3,1)。这样PyTorch会自动将这个维度广播到和a的最后一维（长度5）一致，完成对应位置的乘法。

有两种简洁的写法：

方法1：使用torch.unsqueeze

import torch

# 初始化张量
a = torch.tensor([[[1, 2, 3, 4, 5],
                   [1, 2, 3, 4, 5],
                   [1, 2, 3, 4, 5]],
                  [[1, 2, 3, 4, 5],
                   [1, 2, 3, 4, 5],
                   [1, 2, 3, 4, 5]]])
b = torch.tensor([[2.0, 1.0, 2.0],[0.5, 1.0, 1.0]])

# 扩展b的维度并相乘
result = a * b.unsqueeze(-1)
print(result)

方法2：使用索引语法[:, :, None]

这是PyTorch中更简洁的维度扩展写法，效果和unsqueeze(-1)完全一致：

result = a * b[:, :, None]
print(result)

输出结果 #

运行上述代码后，会得到你期望的结果：

tensor([[[ 2.0000,  4.0000,  6.0000,  8.0000, 10.0000],
         [ 1.0000,  2.0000,  3.0000,  4.0000,  5.0000],
         [ 2.0000,  4.0000,  6.0000,  8.0000, 10.0000]],

        [[ 0.5000,  1.0000,  1.5000,  2.0000,  2.5000],
         [ 1.0000,  2.0000,  3.0000,  4.0000,  5.0000],
         [ 1.0000,  2.0000,  3.0000,  4.0000,  5.0000]]])

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