pytorch 广播机制 pytorch broadcast

本文参考的资料有PYTORCH BROADCASTING SEMANTICS¹、NUMPY BROADCASTING²，搭配一些个人理解。

什么是广播机制？

如果一个Pytorch运算支持广播的话，那么就意味着传给这个运算的参数会被自动扩张成相同的size，在不复制数据的情况下就能行。
广播机制实际上就是在运算过程中，去处理两个形状不同向量的一种手段²。
Numpy通过广播机制³，可以让循环在C中而不是Python中进行。整个过程可以做到避免无用的复制，达到更高效的运算。

两个张量可广播的前提

前提1: 两个张量都至少有一个维度。
像下面这种情况下就不行，因为x不满足这个条件。

x=torch.empty((0,))
y=torch.empty(2,2)
 

前提2: 按顺序看两个张量的每一个维度，x和y每个对应着的两个维度都需要能够匹配上。什么情况下算是匹配上了？满足下面的条件就可以：

if 这两个维度的大小相等
elif 某个维度 一个张量有，一个张量没有
elif 某个维度 一个张量有，一个张量也有但大小是1
 
这样说也许有点抽象，下面这个例子我用颜色标出了x、y每两个对应维度是如何分别符合了上面的条件：

x=torch.empty(5,3,4,1)
y=torch.empty( 3,1,1)

当两个张量满足可广播前提后，具体如何广播？

1. 首先第一步，上面说的蓝色情况会被统一成橙色情况。
   举例：
   统一前:
   x=torch.empty(5,3,4,1)
   y=torch.empty( 3,1,1)
   统一后：
   x=torch.empty(5,3,4,1)
   y=torch.empty(1,3,1,1)
   也就是说蓝色情况的那个缺失维度会被加上。
2. 第二步，现在就只剩下橙色情况需要被解决了。x、y对应维度不相等的时候，size为1的维度会被广播得和对应维度一样大。比如y中0维的1会变成5，y中2维的1会变成4。

实例

“相加”这个运算是支持广播的。那么使用两个满足可广播条件的张量进行实验，就有：

>>> x=torch.empty(5,1,4,1)
>>> y=torch.empty(  3,1,1)
>>> (x+y).size()
torch.Size([5, 3, 4, 1])

可以看到广播的过程是自动完成的。
 

从空间上理解广播机制

举两个例子(图片来源见注解²)：
例1：

>>> a = np.array([1.0, 2.0, 3.0])
>>> b = 2.0
>>> a * b
array([ 2.,  4.,  6.])

解析：这个例子很简单，b被自动广播得和a一样大，完成了加运算。

例2：

>>> a = np.array([0.0, 10.0, 20.0, 30.0])
>>> b = np.array([1.0, 2.0, 3.0])
>>> a[:, np.newaxis] + b
array([[  1.,   2.,   3.],
       [ 11.,  12.,  13.],
       [ 21.,  22.,  23.],
       [ 31.,  32.,  33.]])

解析：这个例子稍微复杂一丢丢。np.newaxis就是给a新添了一个维度，于是加运算的时候：

a的形状是(4,1)

b的形状是(3)

那么如果a和b要匹配上，第一步给b新添一个维度，我们有：
a的形状是(4,1)
b的形状是(1,3)

第二步二者各自把为1的维度进行广播，就如图中那样，最后运算完成。

---

1. PYTORCH BROADCASTING SEMANTICS ↩︎
2. NUMPY BROADCASTING ↩︎ ↩︎ ↩︎
3. 作者注：PYTORCH BROADCASTING SEMANTICS中开头便说，“Many PyTorch operations support NumPy’s broadcasting semantics.”，Pytorch和Numpy的广播机制应当是一样的。 ↩︎