gm_pooling.py

from torch.nn.parameter import Parameter
import torch.nn.functional as F
import torch.nn as nn
import torch
def gem(x, p=3, eps=1e-6):
    return F.avg_pool2d(x.clamp(min=eps).pow(p), (x.size(-2), x.size(-1))).pow(1./p)
class GeM(nn.Module):
    def __init__(self, p=3, eps=1e-6):
        super(GeM,self).__init__()
        self.p = Parameter(torch.ones(1)*p)
        self.eps = eps
    def forward(self, x):
        return gem(x, p=self.p, eps=self.eps)
    def __repr__(self):
        return self.__class__.__name__ + '(' + 'p=' + '{:.4f}'.format(self.p.data.tolist()[0]) + ', ' + 'eps=' + str(self.eps) + ')'

class MyDropOut(nn.Module):
    def __init__(self,p=0.5):
        super(MyDropOut, self).__init__()
        self.p = p
        self.dropout1 = nn.Dropout(self.p)
        self.dropout2 = nn.Dropout(self.p)

    def forward(self, x):
        x = self.dropout1(x)
        return x

    def __repr__(self):
            return self.__class__.__name__ +"dropout"