Numpy、Pytorch中的broadcasting

寫在前面

自己一直都不清楚numpy、pytorch里面不同維數(shù)的向量之間的element wise的計算究竟是按照什么規(guī)則來確認(rèn)維數(shù)匹配和不匹配的情況的，比如

>>> b = np.ones((4,5))
>>> a = np.arange(5)
>>> c = a + b
>>> c.shape
(4, 5)
>>> c
array([[1., 2., 3., 4., 5.],
? ? ? ?[1., 2., 3., 4., 5.],
? ? ? ?[1., 2., 3., 4., 5.],
? ? ? ?[1., 2., 3., 4., 5.]])

上面這種情況就會自動讓a和b的維數(shù)匹配，a加到了b的每一行上

>>> b = np.ones((5,4))
>>> a = np.arange(5)
>>> c = a + b
Traceback (most recent call last):
? File "<stdin>", line 1, in <module>
ValueError: operands could not be broadcast together with shapes (5,) (5,4)

這種情況就無法匹配，此時我們希望的是a能自動加到b的每一列上，但結(jié)果看來好像不行

雖然一直存在這種疑惑，但因為平時遇到的各種運算都比較簡單，遇到這種不是直接匹配的array的加法第一直覺就是去console里面試一試，報錯就換個姿勢再試一試，總歸問題可以快速地解決，但是最近在寫模型的時候，遇到了繞不過去的問題，所以去查了文檔，本文就以解決那個問題為目標(biāo)，來解釋清楚pytorch(numpy也是一樣)中的broadcasting semantics的問題

問題描述

我有一個數(shù)據(jù)Tensor，維數(shù)是64 × 2048 64\times204864×2048，現(xiàn)在我想通過對這64 6464個2048 20482048維的向量做attention（也就是做一個加權(quán)和）來得到一個2048 20482048維的向量，因為模型的需要，我需要用五組不同的權(quán)值向量來計算出五個不同的加權(quán)結(jié)果，也就是我的計算結(jié)果應(yīng)該是一個5 × 2048 5\times 20485×2048維的向量，因為在64 6464個向量上加權(quán)，所以一組權(quán)值向量是64 6464維，五組就是5 × 64 5\times 645×64維

嘗試解決

現(xiàn)在我手頭上有兩個Tensor，一個是數(shù)據(jù)Tensor（64 × 2048 64\times 204864×2048）另一個是權(quán)值Tensor(5 × 64 5\times 645×64)，我GAN!直到我寫到了這里，我才發(fā)現(xiàn)這不是一個矩陣乘法就能解決的問題嘛+_+，當(dāng)然，我想給自己正名，這里我簡化了一下問題所以才發(fā)現(xiàn)原來這么容易就解決了，而原來我在寫代碼的時候因為還要考慮batch_size等問題才云里霧里不知道咋辦，還好當(dāng)時沒想出來，所以去查了文檔發(fā)現(xiàn)了新的東西，然后寫文章的時候想到也算是完滿了（不然也不會發(fā)現(xiàn)自己好澇）

以上都是題外話，現(xiàn)在，我們還是考慮用愚蠢的element wise的方法來解決，好在現(xiàn)在有兩種方法可以解決問題，所以我們可以用來相互檢驗一下，element wise的解決方法就是，我希望這5個64維的權(quán)值向量分別和這64個2048維的向量進(jìn)行element wise的乘法，也就是第一個64維權(quán)值向量先對64個2048維向量加權(quán)得到一個2048維的向量，然后第二個64維權(quán)值向量先對64個2048維向量加權(quán)得到一個2048維的向量…，以此類推總共五個，最終得到五個64 × 2048 64×204864×2048維的向量，然后求和得到最后的5 × 2048 5×20485×2048維的向量

那么按照平常的習(xí)慣，我就去先試試pytorch能不能直接地理解我的想法

>>> import torch
>>> bs = 10 # batch_size
>>> x = torch.randn(bs,64,2048)
>>> att = torch.randn(5,64)
>>> out = att * x
Traceback (most recent call last):
? File "<stdin>", line 1, in <module>
RuntimeError: The size of tensor a (64) must match the size of tensor b (2048) at non-singleton dimension 2

直接乘不行，因為維數(shù)是不匹配的，那怎樣的維數(shù)才算匹配呢？

BROADCASTING SEMANTICS

以下內(nèi)容主要來源于自官方文檔

很多pytorch的運算是支持broadcasting semantics的，而簡單來說，如果運算支持broadcast，則參與運算的Tensor會自動進(jìn)行擴展來使得運算符左右的Tensor維數(shù)匹配，而無需人手動地去拷貝其中的某個Tensor，這就類似于我們開頭的那個例子

>>> b = np.ones((4,5))
>>> a = np.arange(5)
>>> c = a + b
>>> c.shape
(4, 5)
>>> c
array([[1., 2., 3., 4., 5.],
? ? ? ?[1., 2., 3., 4., 5.],
? ? ? ?[1., 2., 3., 4., 5.],
? ? ? ?[1., 2., 3., 4., 5.]])

我們無需讓a的維數(shù)和b一樣，因為numpy自動幫我們做了

這里的另一個重要的概念是broadcastable，如果兩個Tensor是broadcastable的，那么就可以對他倆使用支持broadcast的運算，比如直接加減乘除

而兩個向量要是broadcast的話，必須滿足以下兩個條件

• 每個tensor至少是一維的
• 兩個tensor的維數(shù)從后往前，對應(yīng)的位置要么是相等的，要么其中一個是1，或者不存在

這是官方的例子解釋

>>> x=torch.empty(5,7,3)
>>> y=torch.empty(5,7,3)
# 相同維數(shù)的tensor一定是broadcastable的

>>> x=torch.empty((0,))
>>> y=torch.empty(2,2)
# 不是broadcastable的，因為每個tensor維數(shù)至少要是1

>>> x=torch.empty(5,3,4,1)
>>> y=torch.empty( ?3,1,1)
# 是broadcastable的，因為從后往前看，一定要注意是從后往前看！
# 第一個維度都是1，相等，滿足第二個條件
# 第二個維度其中有一個是1，滿足第二個條件
# 第三個維度都是3，相等，滿足第二個條件
# 第四個維度其中有一個不存在，滿足第二個條件

# 但是
>>> x=torch.empty(5,2,4,1)
>>> y=torch.empty( ?3,1,1)
# 不是broadcastable的，因為從后往前看第三個維度是不match的 2!=3，且都不是1

如果x和y是broadcastable的，那么結(jié)果的tensor的size按照如下的規(guī)則計算

• 如果兩者的維度不一樣，那么就自動增加1維（也就是unsqueeze）
• 對于結(jié)果的每個維度，它取x和y在那一維上的最大值

官方的例子

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

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

>>> x=torch.empty(5,2,4,1)
>>> y=torch.empty(3,1,1)
>>> (x+y).size()
RuntimeError: The size of tensor a (2) must match the size of tensor b (3) at non-singleton dimension 1

此外，關(guān)于broadcast導(dǎo)致的就地（in-place）操作和梯度運算的兼容性等問題，可以自行參考官方文檔

解決問題

上面我們看到，要想兩個Tensor支持element wise的運算，需要它們是broadcastable的，而要想它們是broadcastable的，就需要它們的維度自后向前逐一匹配，回到我們原來的問題中，我們有兩個Tensor x(64 × 2048) att(5 × 64)，為了讓它們broadcastable，我們只需要

>>> import torch
>>> bs = 10 # batch_size
>>> x = torch.randn(bs,64,2048)
>>> att = torch.randn(5,64)
>>> x = x.unsqueeze(1)
>>> att = att.view(1,*att.shape,1)
>>> x.shape
torch.Size([10, 1, 64, 2048])
>>> att.shape
torch.Size([1, 5, 64, 1])
>>> out = x * att
>>> out.shape
torch.Size([10, 5, 64, 2048])

最后我們來驗證兩種方法是否結(jié)果相同

>>> import torch
>>> bs = 10?
>>> x = torch.randn(bs,64,2048)
>>> att = torch.randn(5,64)
>>> out1 = torch.matmul(att,x) ?# 直接矩陣相乘
>>> out.shape
torch.Size([10, 5, 2048])

>>> x = x.unsqueeze(1)
>>> att = att.view(1,*att.shape,1)
>>> out2 = x * att ?# element wise的方法
>>> out2 = out2.sum(dim=2)

>>> test = torch.sum((out1-out2)<0.00001) ?# 浮點數(shù)有微小的誤差
>>> test
tensor(102400)
>>> out1.numel() ?# 最后表明兩個out向量是相等的
102400

Reference

[1] https://docs.scipy.org/doc/numpy/user/basics.broadcasting.html#module-numpy.doc.broadcasting

[2] https://pytorch.org/docs/stable/notes/broadcasting.html#broadcasting-semantics

總結(jié)

以上為個人經(jīng)驗，希望能給大家一個參考，也希望大家多多支持腳本之家。

最新評論