PyTorch中如何不使用循环实现沿新维度的张量乘法

核心逻辑是利用张量的广播机制，通过给低维张量插入扩展维度，让两个张量的维度匹配后直接做元素级乘法，不需要编写for循环。

通用实现思路 #

已知t1形状为[a, b]，t2形状为[c]，只需要给t1的最后一维插入一个扩展维度，让t1形状变为[a, b, 1]，形状为[c]的t2会自动广播为[1, 1, c]，二者做元素乘法后得到的结果就是形状为[a, b, c]的目标张量t3，完全符合你给出的运算规则。

各主流框架代码实现 #

NumPy 实现 #

import numpy as np

# 假设t1、t2为已定义的numpy数组
t3 = t1[..., np.newaxis] * t2
# 等价写法：t3 = np.expand_dims(t1, axis=-1) * t2

PyTorch 实现 #

import torch

# 假设t1、t2为已定义的torch张量
t3 = t1.unsqueeze(-1) * t2
# 等价写法：t3 = t1[..., None] * t2

TensorFlow 实现 #

import tensorflow as tf

# 假设t1、t2为已定义的tf张量
t3 = tf.expand_dims(t1, axis=-1) * t2

结果验证示例 #

我们用一组实际数值验证运算正确性：
设定a=2、b=3、c=2：

• t1 = [[1,2,3],[4,5,6]]，形状[2,3]
• t2 = [10, 20]，形状[2]

运算得到的t3结果如下，和预期规则完全匹配：

# 最后一维第0个切片
[:, :, 0] = [[10, 20, 30], [40, 50, 60]]
# 最后一维第1个切片
[:, :, 1] = [[20, 40, 60], [80, 100, 120]]

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