torch.nn.AdaptiveAvgPool分析

本文分析pytorch AdaptiveAvgPool算子的计算机制,主要是指确定stride和kernel_size的方法。

1. AdaptiveAvgPool1d

目前网上常见的计算1d的公式是与pytorch实现方式不一样的torch.nn.AdaptiveAvgPool分析_第1张图片
通过阅读源码,可以发现,pytorch中adaptive_avg_pool1d实现方式实际上是将一维tensor扩展为二维tensor,之后调用adaptive_avg_pool2d来实现的。

2. AdaptiveAvgPool2d

参考源码可知,pytorch AdaptiveAvgPool2d中stride和kernel_size并不是通过计算output_shape和input_shape来确定的,实际上是通过如下公式,在每一个output点上分别计算
对于output中的某个点i:

//伪代码,说明数学关系
int start_index(int a, int b, int c) //计算下采样起始位置
{
    return floor((float)(a * c) / b);
}
int end_index(int a, int b, int c) //计算下采样结束位置
{
    return ceil((float)((a + 1) * c) / b);
}

//output第i个点在H轴上的下采样(pooling)范围
istartH = start_index(ih, osizeH, isizeH);
iendH  = end_index(ih, osizeH, isizeH);
kH = iendH - istartH;
//output第i个点在W轴上的下采样(pooling)范围
istartW = start_index(iw, osizeW, isizeW);
iendW  = end_index(iw, osizeW, isizeW);
kW = iendW - istartW;

output[i].H = sum(input[istartH:iendH]) / kH  //H轴上的下采样
output[i].W = sum(input[istartW:iendW]) / kW  //W轴上的下采样

3.分析

这样做的好处是不用考虑output和input的shape,如果采用统一计算s,k的公式,当遇到无法整除的情况时,会进行较为复杂的特殊情况处理,采用这种一一确定采样点的方法适配性更强。

你可能感兴趣的:(tvm,pytorch)