前言
最近用 mindspore 復現 Fast SCNN 網絡的時候,裡面用到了一個自適應平均池化算子 「nn.AdaptiveAvgPool2d」,但是 mindspore 目前的版本還沒有提供對應的算子供開發者使用,所以筆者查閱了一部分資料,了解了其計算原理。實驗結果表明,下述方法可以替代 AdaptiveAvgPool2d 算子,可以此為基礎遷移至其他 AI 框架(比如mindspore、tensorflow)中。
一、AdaptiveAvgPool2d
AdaptiveAvgPool2d 的功能簡單來講就是,開發者只需要傳入「待處理數據」 和 「目標大小」,該算子會自動計算池化操作時的kernel_size和stride等數據,使得輸出結果的shape為「目標大小」。
但其實上述理解是不完全正確的,甚至是錯誤的。基於此理解可以稍微理解其功能,但是如果想對其進行復現,則會完全陷入誤區。
目前普遍存在的一種復現方式是,既然我們知道普通池化操作的計算過程是:
已知池化層的kernel_size、padding、stride 以及輸入張量的大小input_size,則輸出張量大小 output_size 為:
output_size =(input_size+2*padding-kernel_size)/stride +1
(此處簡化了計算,如果輸入張量的 column 值和 row 值不等,則分別計算)
那麼我們就想辦法通過 input_size 和 output_size 反推出 kernel_size、stride 等數值就好了,反向回去計算一定可以得到我們需要的數據。
但其實這種方法僅合其形,不對其意。我們只能使得輸出張量是我們需要的「目標大小」,其內部數值卻和 「nn.AdaptiveAvgPool2d」 的計算結果有不小差異,究其原因,在於出發點的錯誤。
二、AdaptiveAvgPool2d 計算原理
筆者翻閱了不少資料,最終在https://discuss.pytorch.org/t/what-is-adaptiveavgpool2d/26897 找到了我需要的內容。其中 Thomas 對AdaptiveAvgPool2d 做了相當準確的解釋,我將其分享的代碼改進為了 NCHW 的模式:
import torch.nn as nn
import torch
def torch_pool(inputs, target_size):
#NCHW
H = target_size[0]
W = target_size[1]
s_p1 = (torch.arange(W, dtype=torch.float32) * (inputs.size(-1) / W)).long()
e_p1 = ((torch.arange(W, dtype=torch.float32)+1) * (inputs.size(-1) / W)).ceil().long()
s_p2 = (torch.arange(H, dtype=torch.float32) * (inputs.size(-2) / H)).long()
e_p2 = ((torch.arange(H, dtype=torch.float32)+1) * (inputs.size(-2) / H)).ceil().long()
pooled2 = []
for i_H in range(H):
pooled = []
for i_W in range(W):
res = torch.mean(inputs[:, :, s_p2[i_H]:e_p2[i_H],s_p1[i_W]:e_p1[i_W]], dim=(-2,-1), keepdim=True)
pooled.append(res)
pooled = torch.cat(pooled, -1)
pooled2.append(pooled)
pooled2 = torch.cat(pooled2,-2)
return pooled2
if __name__ == '__mAIn__':
data = [[[[2,3,4,5,6,9,7,8]
,[2,3,4,5,6,9,7,8]
,[2,3,4,5,6,9,7,8]
,[2,3,4,5,6,9,7,8]
,[2,3,4,5,6,9,7,8]
,[2,3,4,5,6,9,7,8]]
,
[[2,3,4,5,6,9,7,8]
,[2,3,4,5,6,9,7,8]
,[2,3,4,5,6,9,7,8]
,[2,3,4,5,6,9,7,8]
,[2,3,4,5,6,9,7,8]
,[2,3,4,5,6,9,7,8]]
]]
inputs = torch.tensor(data,dtype=torch.float32)
print(inputs)
print(inputs.size())
print("*********************************")
avgpool1 = torch_pool(inputs, (1,3))
avgpool2 = torch_pool(inputs, (2,3))
avgpool3 = torch_pool(inputs, (3,3))
avgpool6 = torch_pool(inputs, (6,5))
print(avgpool1)
print("*********************************")
print(avgpool2)
print("*********************************")
print(avgpool3)
print("*********************************")
print(avgpool6)
計算結果:
tensor([[[[2., 3., 4., 5., 6., 9., 7., 8.],
[2., 3., 4., 5., 6., 9., 7., 8.],
[2., 3., 4., 5., 6., 9., 7., 8.],
[2., 3., 4., 5., 6., 9., 7., 8.],
[2., 3., 4., 5., 6., 9., 7., 8.],
[2., 3., 4., 5., 6., 9., 7., 8.]],
[[2., 3., 4., 5., 6., 9., 7., 8.],
[2., 3., 4., 5., 6., 9., 7., 8.],
[2., 3., 4., 5., 6., 9., 7., 8.],
[2., 3., 4., 5., 6., 9., 7., 8.],
[2., 3., 4., 5., 6., 9., 7., 8.],
[2., 3., 4., 5., 6., 9., 7., 8.]]]])
torch.Size([1, 2, 6, 8])
*********************************
tensor([[[[3., 6., 8.]],
[[3., 6., 8.]]]])
*********************************
tensor([[[[3., 6., 8.],
[3., 6., 8.]],
[[3., 6., 8.],
[3., 6., 8.]]]])
*********************************
tensor([[[[3., 6., 8.],
[3., 6., 8.],
[3., 6., 8.]],
[[3., 6., 8.],
[3., 6., 8.],
[3., 6., 8.]]]])
*********************************
tensor([[[[2.5000, 4.0000, 5.5000, 7.3333, 7.5000],
[2.5000, 4.0000, 5.5000, 7.3333, 7.5000],
[2.5000, 4.0000, 5.5000, 7.3333, 7.5000],
[2.5000, 4.0000, 5.5000, 7.3333, 7.5000],
[2.5000, 4.0000, 5.5000, 7.3333, 7.5000],
[2.5000, 4.0000, 5.5000, 7.3333, 7.5000]],
[[2.5000, 4.0000, 5.5000, 7.3333, 7.5000],
[2.5000, 4.0000, 5.5000, 7.3333, 7.5000],
[2.5000, 4.0000, 5.5000, 7.3333, 7.5000],
[2.5000, 4.0000, 5.5000, 7.3333, 7.5000],
[2.5000, 4.0000, 5.5000, 7.3333, 7.5000],
[2.5000, 4.0000, 5.5000, 7.3333, 7.5000]]]])
與 pytorch 的 nn.AdaptiveAvgPool2d 算子進行對比驗證:
import torch.nn as nn
import torch
if __name__ == '__main__':
data = [[[[2,3,4,5,6,9,7,8]
,[2,3,4,5,6,9,7,8]
,[2,3,4,5,6,9,7,8]
,[2,3,4,5,6,9,7,8]
,[2,3,4,5,6,9,7,8]
,[2,3,4,5,6,9,7,8]]
,
[[2,3,4,5,6,9,7,8]
,[2,3,4,5,6,9,7,8]
,[2,3,4,5,6,9,7,8]
,[2,3,4,5,6,9,7,8]
,[2,3,4,5,6,9,7,8]
,[2,3,4,5,6,9,7,8]]
]]
x = torch.tensor(data,dtype=torch.float32)
print(x)
print(x.size())
print("*********************************")
avgpool1 = nn.AdaptiveAvgPool2d((1,3))
avgpool2 = nn.AdaptiveAvgPool2d((2,3))
avgpool3 = nn.AdaptiveAvgPool2d((3,3))
avgpool6 = nn.AdaptiveAvgPool2d((6,5))
print(avgpool1(x))
print("*********************************")
print(avgpool2(x))
print("*********************************")
print(avgpool3(x))
print("*********************************")
print(avgpool6(x))
計算結果:
tensor([[[[2., 3., 4., 5., 6., 9., 7., 8.],
[2., 3., 4., 5., 6., 9., 7., 8.],
[2., 3., 4., 5., 6., 9., 7., 8.],
[2., 3., 4., 5., 6., 9., 7., 8.],
[2., 3., 4., 5., 6., 9., 7., 8.],
[2., 3., 4., 5., 6., 9., 7., 8.]],
[[2., 3., 4., 5., 6., 9., 7., 8.],
[2., 3., 4., 5., 6., 9., 7., 8.],
[2., 3., 4., 5., 6., 9., 7., 8.],
[2., 3., 4., 5., 6., 9., 7., 8.],
[2., 3., 4., 5., 6., 9., 7., 8.],
[2., 3., 4., 5., 6., 9., 7., 8.]]]])
torch.Size([1, 2, 6, 8])
*********************************
tensor([[[[3., 6., 8.]],
[[3., 6., 8.]]]])
*********************************
tensor([[[[3., 6., 8.],
[3., 6., 8.]],
[[3., 6., 8.],
[3., 6., 8.]]]])
*********************************
tensor([[[[3., 6., 8.],
[3., 6., 8.],
[3., 6., 8.]],
[[3., 6., 8.],
[3., 6., 8.],
[3., 6., 8.]]]])
*********************************
tensor([[[[2.5000, 4.0000, 5.5000, 7.3333, 7.5000],
[2.5000, 4.0000, 5.5000, 7.3333, 7.5000],
[2.5000, 4.0000, 5.5000, 7.3333, 7.5000],
[2.5000, 4.0000, 5.5000, 7.3333, 7.5000],
[2.5000, 4.0000, 5.5000, 7.3333, 7.5000],
[2.5000, 4.0000, 5.5000, 7.3333, 7.5000]],
[[2.5000, 4.0000, 5.5000, 7.3333, 7.5000],
[2.5000, 4.0000, 5.5000, 7.3333, 7.5000],
[2.5000, 4.0000, 5.5000, 7.3333, 7.5000],
[2.5000, 4.0000, 5.5000, 7.3333, 7.5000],
[2.5000, 4.0000, 5.5000, 7.3333, 7.5000],
[2.5000, 4.0000, 5.5000, 7.3333, 7.5000]]]])
可以發現,無論是輸出 shape 還是 Tensor 內部數值,二者都是一樣的。
而且無論改變 NCHW 的哪一部分內容,擴充 N、C、H、W 的任意一維,最後的計算結果都是保持一致的。
三、AdaptiveAvgPool2d 計算原理詳解
待補充。
至此,探究明白了 「nn.AdaptiveAvgPool2d」 的內部計算原理。並可以此為基礎遷移至其他 AI 框架中。
四、mindspore版 AdaptiveAvgPool2d
觀察上述代碼,如果要重寫成mindspore版的代碼,我們只需要替換掉『torch.arange』、『torch.mean』、『torch.cat』這三個主要算子,以及添加一個取整操作,在mindspore中就是ops.ReduceMean(keep_dims=True)、P.Concat(axis=-1)等算子,只要做對應替換就可以了。
但我之前在重寫 res = torch.mean(inputs[:, :, s_p2[i_H]:e_p2[i_H],s_p1[i_W]:e_p1[i_W]], dim=(-2,-1),
keepdim=True)這一句時,發現mindspore對『變量下標』做切片操作時會發生異常,不太清楚是不是我的用法有問題。不過我寫了一個臨時的版本,比如如果要將NCx32x64的數據池化成NCx6x6大小,我們可以提前計算出需要切片的下標,就可以得到這樣一版可用的代碼了:
def _AvgPool2d6x6(self,x):
s_p1 = [ 0, 10, 21, 32, 42, 53]
e_p1 = [11, 22, 32, 43, 54, 64]
s_p2 = [ 0, 5, 10, 16, 21, 26]
e_p2 = [ 6, 11, 16, 22, 27, 32]
pooled2 = []
for i_H in range(6):
pooled = []
for i_W in range(6):
res = self.reduceMean(x[:, :, s_p2[i_H]:e_p2[i_H],s_p1[i_W]:e_p1[i_W]], (-2,-1))
pooled.append(res)
pooled = self.concat1((pooled[0],pooled[1],pooled[2],pooled[3],pooled[4],pooled[5]))
pooled2.append(pooled)
pooled2 = self.concat2((pooled2[0],pooled2[1],pooled2[2],pooled2[3],pooled2[4],pooled2[5]))
return pooled2
原載於http://luxuff.cn ,作者還是我自己。
轉載請註明出處。