PyTorch 广播机制

2024-09-06 17:55:29

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

最新文章

热门文章