PyTorch 广播机制

PyTorch 广播机制

定义

PyTorch的tensor参数可以自动扩展其大小。一般的是小一点的会变大，来满足运算需求。

规则

满足一下情况的tensor是可以广播的。

• 至少有一个维度
• 两个tensor维度相等
• 维度不等，其中一个为1
• 维度不等，其中一个维度不存在

计算过程规则。

• 维度不同，小维度的增加维度
• 每个维度，计算结果取大的
• 扩展维度是对数值进行复制

例子

# 广播机制
a = torch.arange(3).reshape((3, 1))

b = torch.arange(2).reshape((1, 2))
print(a + b)

输出结果

tensor([[0],
        [1],
        [2]])
tensor([[0, 1]])
-----------------------
tensor([[0, 1],
        [1, 2],
        [2, 3]])