pytorch torch.nn.AdaptiveAvgPool2d()自适应平均池化函数详解

PyTorch torch.nn.AdaptiveAvgPool2d() 自适应池化函数详解

池化操作简介

在深度学习的卷积神经网络（Convolutional Neural Network，CNN）中，池化操作是常用的一种非线性操作，用于缩小特征图尺寸和提取主要特征。

普通池化操作

普通池化操作，也称为固定池化（Fixed Pooling），是指一种对特征图按照固定大小裁剪并执行固定步幅下采样的操作。普通池化操作一般有 Max Pooling、Mean Pooling 两种，其中 Max Pooling 取所有像素点中最大值，而 Mean Pooling 取像素点的平均值。

自适应池化操作

相比于固定池化操作，自适应池化操作则可动态地针对特征图的大小进行调整，以适应不同大小的输入特征图和不同的网络结构。

PyTorch 中的自适应池化函数 torch.nn.AdaptiveAvgPool2d() 即可实现自适应平均池化操作。

torch.nn.AdaptiveAvgPool2d()

介绍

torch.nn.AdaptiveAvgPool2d(output_size)

自适应平均池化操作，它会对输入的 feature map 进行自适应的裁剪与池化，选择输出大小为 output_size 的 feature map。

具体地，该函数的输入为一个四维的 Tensor，维度分别为 (batch_size, C, H, W)，其中 batch_size 表示输入数据的 batch 大小，C 表示输入数据的通道数，H 和 W 分别为输入数据的高度和宽度。

该函数的输出也是一个四维的 Tensor，其维度同样为 (batch_size, C, H, W)，其中 H 和 W 为 output_size所指定的值。

参数

• output_size (tuple[int]) – 输出的特征图的大小，格式为 (height, width) ，也可以是单个 int 值，表示输出特征图高度和宽度的相同长度。

用法示例1

我们来看一个使用 AdaptiveAvgPool2d() 的示例，这个示例代码用于实现任意大小输入图像的分类模型。

import torch.nn as nn
import torch

class CNN(nn.Module):
    def __init__(self):
        super(CNN, self).__init__()
        self.conv = nn.Sequential(
            nn.Conv2d(3, 16, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(16),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=2, stride=2),
            nn.Conv2d(16, 32, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(32),
            nn.ReLU(inplace=True),
            nn.AdaptiveAvgPool2d((1,1)),
        )
        self.fc = nn.Linear(32, 10)

    def forward(self, x):
        x = self.conv(x)
        x = x.view(x.size(0), -1)
        x = self.fc(x)
        return x

net = CNN()
x = torch.randn(1, 3, 32, 32)
y = net(x)
print(y.shape)

在上面的示例代码中，首先使用了一组卷积、池化来提取特征，最后使用了自适应池化层将得到的特征图转换成向量，再通过全连接层完成最终的分类任务。由于自适应池化层的输入大小可以是任意的，所以无论输入图像的大小是多少，都可以适应。

用法示例2

下面我们再来看一个使用 AdaptiveAvgPool2d() 的示例，这个示例代码用于提取图像的全局特征。

import torch
import torch.nn as nn

class Flatten(nn.Module):
    def forward(self, x):
        return x.view(x.shape[0], -1)

class Model(nn.Module):
    def __init__(self, input_shape=(3, 64, 64), output_dim=10):
        super().__init__()

        self.features = nn.Sequential(
            nn.Conv2d(input_shape[0], 32, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.BatchNorm2d(32),
            nn.Conv2d(32, 64, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.BatchNorm2d(64),
            nn.Conv2d(64, 128, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.BatchNorm2d(128),
            nn.Conv2d(128, 256, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.BatchNorm2d(256),
        )

        self.avgpool = nn.AdaptiveAvgPool2d(output_size=1)
        self.flatten = Flatten()
        self.fc = nn.Linear(256, output_dim)

    def forward(self, x):
        z = self.features(x)
        z = self.avgpool(z)
        z = self.flatten(z)
        logits = self.fc(z)
        return logits

net = Model()
x = torch.randn(1, 3, 64, 64)
y = net(x)
print(y.shape)

在上面的示例代码中，首先使用一组卷积层提取图像的特征，随后使用 nn.AdaptiveAvgPool2d(output_size=1) 自适应平均池化层将其转换为全局特征。最后通过全连接层完成最终的分类任务。

总结

torch.nn.AdaptiveAvgPool2d() 函数是 PyTorch 提供的一个自适应平均池化函数，可以根据输入数据的大小动态地进行特征图的裁剪和池化，适用于各种大小的输入特征图和不同的网络结构。可以用于构建各种基于卷积神经网络模型的分类、分割、检测等任务。