理解pytorch的广播语义

news2025/1/16 21:48:42

什么是广播运算

广播的条件

示例

示例1

示例2

示例3 补1

示例4 原位运算

示例5 参与广播运算的两个tensor，必须是从右向左对齐

总结规律

两个tensor可以做广播运算的条件：

两个可以互相广播的tensor运算的步骤：

例子：

什么是广播运算

广播发生的场景很广泛，tensor加法、乘法等都适用广播语义。广播的意思就是，在某些条件下，两个形状不同的tensor仍可以完成运算。

广播的条件

两个tensor可以相互广播的条件是：

1 每一个tensor至少有一个维度(torch.empty((0,))就是一个没有维度的tensor，见下面的例子)

2 从最后一个维度（下面解释何为最后一个维度）开始，逐一比较两个tensor的各个维度，这两个被比较的维度，要么相等，要么有一个是1，要么有一个不存在。

何为最后一个维度？看下面的例子

X 是一个2行3列的tensor。从它最外面那一层[]括号来看，里面有2个元素：

[1,2,3]和[4,5,6]

所以，X的“首维度”就是从[1,2,3],[4,5,6]--->这个方向

再往里一层，是1，2，3的数列，也可以说是4,5,6的数列。这就是X的下一个维度。由于X只有两个维度，所以1,2,3(4,5,6)数列的维度就是x的尾维度，或者说是最后一个维度。

由于x.shape=([2,3])，可见，有两个元素的维度([1,2,3],[4,5,6]这个维度)排在左边，有3个元素的维度(1,2,3（或者说，4,5,6）这个维度)排在右面。故，“首维度”排在shape的最左边，“尾维度”排在最右边

示例

示例1

准备两个tensor

X形状是5,7,3所有元素都是0

Y形状是5,7,3 所有元素都是1

运算结果：

可见，add_操作后，x的每个元素都加1.

示例2

X没有维度，Y是2行2列

可见，两者不能相加

示例3 补1

X是一个3x2x2的tensor，而y是一个只有两个元素的矢量。

这时有人会有疑问：x的维度有俩个“2”，y的“2”应该对应哪一个？

根据官网Broadcasting semantics — PyTorch 2.2 documentation的说法“If the number of dimensions of x and y are not equal, prepend 1 to the dimensions of the tensor with fewer dimensions to make them equal length.”---把1插在维度较少的tensor的前面。说明pytorch是尽量保持尾部维度对齐的，也就是下面的对齐：

3 2 2 x

2 y

然后在y的前面插1：

3 2 2

1 1 2

y前面的维度补了两个“1”，每一次补1，y的外面加一层[]。所以补两个1之后，y变成[ [ [20,30] ] ]

对齐后，沿着第一维度计算：

将x视为一个3元矢量，矢量的每一个元素是一个2x2 tensor；将y视为一个1元矢量，矢量唯一的元素是一个1x2的tensor[[20,30]]。接下来，让这个3元矢量的每一个元素与这个1x2的tensor相加。

把这个2x2的tensor再看成一个二元矢量，每个元素又都是一个2元矢量。而把1x2的tensor看成只有一个元素的矢量，且元素本身又是一个2元矢量[20,30]。于是2x2的tensor加1x2的tensor又可以分解成两个二元矢量分别加一个二元矢量[20,30]。举例看[[1,2], [3,4]] + [20,30]，就是[1,2] + [20,30] 和[3,4] + [20,30].

重复上述操作给[[5,6],[7,8]]和[[9,10],[11,12]]，所以就有了最后的结果：x的每一个元素都加上了[20,30]

示例4 原位运算

所谓原位运算，指的是运算返回值保存在某一个输入变量中，而不是保存在新的变量里。在Karpathy的视频教程中提到，P /= … 就是一个原位运算。示例2的add_操作也是一个原位运算。

如下图所示，x.add_(y)成立，但是y.add_(x)不成立：

原因是，x与y相加时，x的维度不需要变化，而y的规模要变化，才能适配x的形状。由于x.add_(y)的结果保存在x，而不是y中，所以y的形状变化是暂时的，运算结束后，就回到原状态。但是反过来则不行：y.add_(x)导致结果保存在y中，导致运算后y的规模和运算前不同，这是不允许的。

示例5 参与广播运算的两个tensor，必须是从右向左对齐

这个例子说明参与广播运算的两个tensor，必须是从右对齐：

从上面的例子x+y失败，可以看出，虽然y(5x2的tensor)可以在最右边补一个1，变为5x2x1，适配x的规模5x2x4，但是广播语义要求参与运算的两个tensor首先把最右边的维度对齐，然后再补充维度。所以x的最右边维度4是无法匹配y的最右边维度2的，故失败。

总结规律

两个tensor可以做广播运算的条件：

1 两个tensor都至少有一个维度；

2 两个tensor的维度个数要么完全一样，那个维度较少的tensor可以把自己缺少的维度补充为1；

3 补齐可以补充多个维度，但是只能发生在所有已有维度的左边，不能插在已有维度之间，也不能出现在已有维度右边。

4 假如运算是原位运算，则保存运算结果的变量的尺寸不应在运算前后发生变化。

两个可以互相广播的tensor运算的步骤：

1 假如两者维度个数、对应维度的尺寸都相同，则直接对应元素做运算，得出结果即可

2 假如两者维度不同，则

2.1 首先让两个变量的最右边维度对齐

2.2 维度较少的那个变量的左边必然缺少维度，缺少几个，就从最左边开始补几个1

2.3 从最左边开始运算（即是说，先处理x1,y1这一对）。变量x = [x1, x2, x3, … xn]和变量y = [y1, y2, y3,….yn]。把[x2,x3,x4…xn]看作一个元素，把[y2,y3,y4,…yn]看作一个元素。这样，x就被看作是一个含有x1个元素的矢量，y被看作是有y1个元素的矢量。根据前面的描述，不难发现，x1与y1要么相等，要么有一个是1.

2.4 假如x1==y1，则只要把各自对应元素相加即可。每个对应元素又是一个[x2,x3,x4...xn]和[y2,3,y4,..yn]，于是计算[x2,x3,x4...xn]+ [y2,3,y4,..yn]。如果对应元素可以直接相加，就返回结果，否则回到2.3

2.5 假如x1!=y1，那么其中必有一个是1。比如说y1==1。那么x+y可以看成是一个x1维矢量加一个标量。矢量加标量，只要把标量加到矢量的每一个元素即可。矢量的每个元素都是一个[x2,x3,x4….xn]，标量是y1这个维度的元素（y1既然等于1，就只有一个元素）。这两者的加法又回到2.3

重复以上步骤，直到最后的维度（也就是最右边的维度）。