PyTorch/Numpy：如何通过广播实现单矩阵减N个元素得到N个矩阵

完全可以通过PyTorch原生的广播机制实现，不需要写任何Python层循环，核心是调整张量c的维度对齐方式，匹配广播规则即可。

实现原理 #

PyTorch广播运算从最右侧维度开始对齐，满足以下任一条件即可触发广播：

• 两个张量对应维度大小相等
• 其中一个张量对应维度大小为1
• 其中一个张量不存在对应维度

之前调整形状失败，本质是没有把c的元素放到独立的新维度上。我们需要的输出形状是(len(c), *image.shape)也就是(2,2,2)，只要把c的形状调整为(len(c), 1, 1)，运算时c后两个长度为1的维度会自动扩展为和image一致的(2,2)，同时image会自动在最前补一个长度为1的批次维度，两者对齐后直接做减法就能得到目标结果。

代码实现 #

import torch

image = torch.tensor([[6, 9], [8.7, 5.5]])
c = torch.tensor([5.7675, 8.8325])

# 方式1：reshape调整c的形状
result = image - c.reshape(-1, 1, 1)

# 方式2：用None索引扩展维度，写法更简洁，和reshape效果完全等价
# result = image - c[:, None, None]

# 输出结果和列表推导式完全一致
print(result)

运行输出：

tensor([[[ 0.2325,  3.2325],
         [ 2.9325, -0.2675]],

        [[-2.8325,  0.1675],
         [-0.1325, -3.3325]]])

实现优势 #

这种写法完全走PyTorch底层向量化计算，没有Python循环开销，在大尺寸张量场景下，运行效率远高于列表推导+torch.stack的写法。

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