PyTorch 的 Pooling 和 UnPooling函数中的 indices 参数：nn.MaxPool2d/nn.MaxUnpool2dF.max_pool2d/F.max_unpool2d

Posted 2023-03-03 ctrl A_ctrl C_ctrl V

这篇博文主要介绍 PyTorch 的 MaxPooling 和 MAxUnPooling 函数中涉及到的 indices 参数。

indices 是“索引”的意思，对于一些结构对称的网络模型，上采样和下采样的结构往往是对称的，我们可以在下采样做 MaxPooling 的时候记录下来最大值所在的位置，当做上采样的时候把最大值还原到其对应的位置，然后其余的位置补 0 。

indices 参数的作用就是保存 MaxPooling 操作时最大值位置的索引。

如下图所示：

PyTorch 的 torch.nn 和 torch.nn.functional 模块中均有实现 Pooling 和 UnPooling 的 api ，它们的作用和效果是完全相同的。

1、nn.MaxPool2d 和 nn.MaxUnpool2d

使用 nn.MaxPool2d 和 nn.MaxUnpool2d 时要先实例化，事实上 nn 模块下面的函数都是如此（需要先实例化），比如 nn.Conv

import torch
from torch import nn


# 使用 nn.MaxPool2d 和 nn.MaxUnpool2d 时要先实例化
pool = nn.MaxPool2d(kernel_size=2, stride=2, return_indices=True)
unpool = nn.MaxUnpool2d(kernel_size=2, stride=2)

data = torch.tensor([[[[1, 2, 3, 4],
                       [5, 6, 7, 8],
                       [9, 10, 11, 12],
                       [13, 14, 15, 16]]]], dtype=torch.float32)
pool_out, indice = pool(data)
unpool_out = unpool(input=pool_out, indices=indice)

print("pool_out = ", pool_out)
print("indice = ", indice)
print("unpool_out = ", unpool_out)

"""
pool_out =  tensor([[[[ 6.,  8.],
          [14., 16.]]]])
indice =  tensor([[[[ 5,  7],
          [13, 15]]]])
unpool_out =  tensor([[[[ 0.,  0.,  0.,  0.],
          [ 0.,  6.,  0.,  8.],
          [ 0.,  0.,  0.,  0.],
          [ 0., 14.,  0., 16.]]]])
"""

2、F.max_pool2d 和 F.max_unpool2d

使用 F.max_pool2d 和 F.max_unpool2d 时不需要实例化，可以直接使用。

import torch
import torch.nn.functional as F


data = torch.tensor([[[[1, 2, 3, 4],
                       [5, 6, 7, 8],
                       [9, 10, 11, 12],
                       [13, 14, 15, 16]]]], dtype=torch.float32)
pool_out, indice = F.max_pool2d(input=data, kernel_size=2, stride=2, return_indices=True)
unpool_out = F.max_unpool2d(input=pool_out, indices=indice, kernel_size=2, stride=2)

print("pool_out = ", pool_out)
print("indice = ", indice)
print("unpool_out = ", unpool_out)

"""
pool_out =  tensor([[[[ 6.,  8.],
          [14., 16.]]]])
indice =  tensor([[[[ 5,  7],
          [13, 15]]]])
unpool_out =  tensor([[[[ 0.,  0.,  0.,  0.],
          [ 0.,  6.,  0.,  8.],
          [ 0.,  0.,  0.,  0.],
          [ 0., 14.,  0., 16.]]]])
"""

可以看到，nn.MaxPool2d / nn.MaxUnpool2d 和 F.max_pool2d / F.max_unpool2d 的作用和输出结果完全相同。

3、使用 Pooling 和 Conv2d 实现上/下采样的区别和产生的影响

使用 Pooling 和 Conv2d 实现上/下采样的区别主要体现在对奇数大小的特征图的处理中，以特征图大小为 65*65为例。

使用 nn.MaxPool2d 和 F.max_pool2d 实现下采样时，得到的特征图大小是 32*32，上采样得到的特征图大小是 64*64
使用 nn.Conv2d 实现下采样时，得到的特征图大小是 33*33；再使用nn.ConvTranspose2d 上采样得到的特征图大小是 66*66

在很多对称的网络结构中（如 UNet、SegNet）,需要对上采样和下采样的对应的特征图进行大小对齐。

若网络中间某个特征图大小是 65 ，不论使用哪种上/下采样策略，得到的特征图大小必然不可能还是64（不可能是奇数）。此时就要考虑64和66的区别了。

（1）如果使用 64 * 64 大小的特征图，则需要进行 padding 得到 65 * 65 的特征图。但padding操作可能会导致在训练和推理过程中的不确定性问题。

（2）如果使用 66 * 66 大小的特征图，则只需要进行切片得到 65 * 65 的特征图。该操作是更靠谱的。

因此，个人建议使用 nn.Conv2d 和其他函数实现采样操作。