PyTorch广播机制

2022/5/12 posts

# 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]])