layers.py

Go to the documentation of this file.

 
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.nn.parameter import Parameter
 
 
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().__init__()
        self.pp = Parameter(torch.ones(1)*p)
        self.epseps = eps
    def forward(self, x):
        return gem(x, p=self.pp, eps=self.epseps)
    def __repr__(self):
        return self.__class__.__name__ + '(' + 'p=' + '{:.4f}'.format(self.pp.data.tolist()[0]) + ', ' + 'eps=' + str(self.epseps) + ')'
 
 
class Flatten(torch.nn.Module):
    def __init__(self):
        super().__init__()
    def forward(self, x):
        assert x.shape[2] == x.shape[3] == 1, f"{x.shape[2]} != {x.shape[3]} != 1"
        return x[:,:,0,0]
 
 
class L2Norm(nn.Module):
    def __init__(self, dim=1):
        super().__init__()
        self.dimdim = dim
    def forward(self, x):
        return F.normalize(x, p=2, dim=self.dimdim)