如何在 PyTorch 中执行总和池



如何在 PyTorch 中执行总和池。具体来说,如果我们有输入(N, C, W_in, H_in)并希望使用特定kernel_sizestride输出(N, C, W_out, H_out),就像nn.Maxpool2d一样?

您可以使用执行平均池化的torch.nn.AvgPool1d(或torch.nn.AvgPool2dtorch.nn.AvgPool3d(-与总和池成正比。如果确实需要求和值,可以将平均输出乘以池化表面。

https://pytorch.org/docs/stable/generated/torch.nn.AvgPool2d.html#torch.nn.AvgPool2d 找到divisor_override。
设置divisor_override=1
,你会得到一个总和池

import torch
input = torch.tensor([[[1,2,3],[3,2,1],[3,4,5]]])
sumpool = torch.nn.AvgPool2d(2, stride=1, divisor_override=1)
sumpool(input)

你会得到

tensor([[[ 8,  8],
[12, 12]]])

扩展本杰明普朗什的答案:
我也需要总和池,它似乎并不直接存在,但它相当于运行一个由 1 组成的权重参数的 conv2d。我认为运行 AvgPool2d 并乘以内核大小乘积会更快。事实证明,不完全是。

前期底线:
使用torch.nn.functional.avg_pool2d及其相关函数并乘以内核大小。

在Jupyter中测试我发现:

(开销(

%%timeit
x = torch.rand([1,1,1000,1000])
>>> 3.49 ms ± 4.72 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
%%timeit
_=F.avg_pool2d(torch.rand([1,1,1000,1000]), [10,10])*10*10
>>> 4.99 ms ± 74.3 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

(所以 1.50 毫秒± 79.0 毫秒( (我发现*10*10只增加了大约 20 微秒(

avePool = nn.AvgPool2d([10, 10], 1, 0)

%%timeit
_=avePool(torch.rand([1,1,1000,1000]))*10*10
>>> 80.9 ms ± 1.57 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

(所以 77.4 毫秒± 1.58 毫秒(

y = torch.ones([1,1,10,10])

%%timeit
_=F.conv2d(torch.rand([1,1,1000,1000]), y)
>>> 14.4 ms ± 421 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

(因此 10.9 毫秒± 426 微秒(

sumPool = nn.Conv2d(1, 1, 10, 1, 0, 1, 1, False)
sumPool.weight = torch.nn.Parameter(y)
%%timeit
_=sumPool(torch.rand([1,1,1000,1000]))
>>> 7.24 ms ± 63.6 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

(所以 3.75 毫秒± 68.3 微秒(

并作为健全性检查。

abs_err = torch.max(torch.abs(avePool(x)*10*10 - sumPool(x)))
magnitude = torch.max(torch.max(avePool(x)*10*10, torch.max(sumPool(x))))
relative_err = abs_err/magnitude
abs_err.item(), magnitude.item(), relative_err.item()
>>> (3.814697265625e-06, 62.89910125732422, 6.064788493631568e-08)

这可能是一个合理的舍入相关错误。

我不知道为什么功能版本比制作专用内核更快,但看起来如果你想制作专用内核,更喜欢 Conv2D 版本,并在创建内核参数期间使用sumPool.weights.requires_grad = Falsewith torch.no_grad():使权重无法训练。这些结果可能会随着内核大小而变化,因此如果您需要加快此部分,请测试您自己的应用程序。如果我错过了什么,请告诉我...

最新更新