StochasticPooling 随机池化的PyTorch实现

关注

原创 2022-03-17 22:58:55 1.7k 阅读 AI 写同款· GEO 优化 ›

随机池化（Stochastic Pooling）前言 CNN中卷积完后有个步骤叫pooling, 在ICLR2013上，作者Zeiler提出了另一种pooling手段(最常见的就是mean-pooling和max-pooling)，叫stochastic pooling。只需要对Feature Map中的元素按照其概率值大小随机选择，元素选中的概率与其数值大小正相关，并非如同max pooling那样直接选取最大值。这种随机池化操作不但最大化地保证了取值的Max，也部分确保不会所有的元素都被选取max值，从而提高了泛化能力。 ... 阅读详情

网上找了好久随机池化的Pytorch代码，都没有太合适的。就自己写了个不太优美的实现。若是在论文中使用该代码，请私信告知我用途。

from torch import nn
import torch
from torch.nn import functional as F

class StochasticPooling(nn.Module):
    def __init__(self):
        super().__init__()
        self.avg = nn.AvgPool2d(2)

    def forward(self, x):
        s = 4 * F.interpolate(F.avg_pool2d(x, 2), scale_factor=2, mode='nearest')
        p = x / s
        if self.training:
            b, c, h, w = p.shape
            o = torch.zeros(b, c, h // 2, w // 2)
            for i in range(h // 2):
                for j in range(w // 2):
                    pij = p[:, :, 2 * i:2 * i + 2, 2 * j:2 * j + 2].reshape(b, c, -1)
                    idx = torch.distributions.Multinomial(1, pij).sample().reshape(b, c, 2, 2)
                    o[:, :, i, j] = x[:, :, 2 * i:2 * i + 2, 2 * j:2 * j + 2][idx == 1].reshape(b, c)
            return o
        else:
            return 4 * F.avg_pool2d(p * x, 2)

需要注意的是，输入特征图的尺寸大小应该为2的整数倍，否则会报错，是由于下面这两行代码引起的

        s = 4 * F.interpolate(F.avg_pool2d(x, 2), scale_factor=2, mode='nearest')
        p = x / s

收录于

算法

python

标签

#pytorch #随机池化 #Stochastic

池化(Pooling)的种类与具体用法——基于Pytorch 池化的种类、原理与基于Pytorch的具体用法阅读详情

StochasticPooling 随机池化的PyTorch实现

相关推荐