归一化EvoNorms

EvoNorms_PyTorch

https://github.com/lonePatient/EvoNorms_PyTorch

原版说精度提升了一个多点，但是内存占用比原来大了很多，也变慢了

  

   

    

     

    
    

     

      import torch
     
    
   

    

     

    
    

     

      import torch.nn as nn
     
    
   

    

     

    
    

     

      from torch.nn import init
     
    
   

    

     

    
    

     

      from torch.nn.parameter import Parameter
     
    
   

    

     

    
    

     
 
     
    
   

    

     

    
    

     

      def instance_std(x, eps=1e-5):
     
    
   

    

     

    
    

     

       N,C,H,W = x.size()
     
    
   

    

     

    
    

     

       x1 = x.reshape(N*C,-1)
     
    
   

    

     

    
    

     

       var = x1.var(dim=-1, keepdim=True)+eps
     
    
   

    

     

    
    

     
 return var.sqrt().reshape(N,C,1,1)
     
    
   

    

     

    
    

     
 
     
    
   

    

     

    
    

     

      def group_std(x, groups, eps = 1e-5):
     
    
   

    

     

    
    

     

       N, C, H, W = x.size()
     
    
   

    

     

    
    

     

       x1 = x.reshape(N,groups,-1)
     
    
   

    

     

    
    

     

       var = (x1.var(dim=-1, keepdim = True)+eps).reshape(N,groups,-1)
     
    
   

    

     

    
    

     
 return (x1 / var.sqrt()).reshape(N,C,H,W)
     
    
   

    

     

    
    

     
 
     
    
   

    

     

    
    

     
 
     
    
   

    

     

    
    

     

      class BatchNorm2dRelu(nn.Module):
     
    
   

    

     

    
    

     
 def __init__(self,in_channels):
     
    
   

    

     

    
    

     

       super(BatchNorm2dRelu,self).__init__()
     
    
   

    

     

    
    

     

       self.layer = nn.Sequential(
     
    
   

    

     

    
    

     

       nn.BatchNorm2d(in_channels),
     
    
   

    

     

    
    

     

       nn.ReLU(inplace=True))
     
    
   

    

     

    
    

     
 def forward(self, x):
     
    
   

    

     

    
    

     

       output = self.layer(x)
     
    
   

    

     

    
    

     
 return output
     
    
   

    

     

    
    

     
 
     
    
   

    

     

    
    

     
 
     
    
   

    

     

    
    

     

      class EvoNorm2dB0(nn.Module):
     
    
   

    

     

    
    

     
 def __init__(self,in_channels,nonlinear=True,momentum=0.9,eps = 1e-5):
     
    
   

    

     

    
    

     

       super(EvoNorm2dB0, self).__init__()
     
    
   

    

     

    
    

     

       self.nonlinear = nonlinear
     
    
   

    

     

    
    

     

       self.momentum = momentum
     
    
   

    

     

    
    

     

       self.eps = eps
     
    
   

    

     

    
    

     

       self.gamma = Parameter(torch.Tensor(1,in_channels,1,1))
     
    
   

    

     

    
    

     

       self.beta = Parameter(torch.Tensor(1,in_channels,1,1))
     
    
   

    

     

    
    

     
 if nonlinear:
     
    
   

    

     

    
    

     

       self.v = Parameter(torch.Tensor(1,in_channels,1,1))
     
    
   

    

     

    
    

     

       self.register_buffer('running_var', torch.ones(1, in_channels, 1, 1))
     
    
   

    

     

    
    

     

       self.reset_parameters()
     
    
   

    

     

    
    

     
 
     
    
   

    

     

    
    

     
 def reset_parameters(self):
     
    
   

    

     

    
    

     

       init.ones_(self.gamma)
     
    
   

    

     

    
    

     

       init.zeros_(self.beta)
     
    
   

    

     

    
    

     
 if self.nonlinear:
     
    
   

    

     

    
    

     

       init.ones_(self.v)
     
    
   

    

     

    
    

     
 
     
    
   

    

     

    
    

     
 def forward(self, x):
     
    
   

    

     

    
    

     

       N, C, H, W = x.size()
     
    
   

    

     

    
    

     
 if self.training:
     
    
   

    

     

    
    

     

       x1 = x.permute(1, 0, 2, 3).reshape(C, -1)
     
    
   

    

     

    
    

     

       var = x1.var(dim=1).reshape(1, C, 1, 1)
     
    
   

    

     

    
    

     

       self.running_var.copy_(self.momentum * self.running_var + (1 - self.momentum) * var)
     
    
   

    

     

    
    

     
 else:
     
    
   

    

     

    
    

     

       var = self.running_var
     
    
   

    

     

    
    

     
 if self.nonlinear:
     
    
   

    

     

    
    

     

       den = torch.max((var+self.eps).sqrt(), self.v * x + instance_std(x))
     
    
   

    

     

    
    

     
 return x / den * self.gamma + self.beta
     
    
   

    

     

    
    

     
 else:
     
    
   

    

     

    
    

     
 return x * self.gamma + self.beta
     
    
   

    

     

    
    

     
 
     
    
   

    

     

    
    

     
 
     
    
   

    

     

    
    

     

      class EvoNorm2dS0(nn.Module):
     
    
   

    

     

    
    

     
 def __init__(self,in_channels,groups=8,nonlinear=True):
     
    
   

    

     

    
    

     

       super(EvoNorm2dS0, self).__init__()
     
    
   

    

     

    
    

     

       self.nonlinear = nonlinear
     
    
   

    

     

    
    

     

       self.groups = groups
     
    
   

    

     

    
    

     

       self.gamma = Parameter(torch.Tensor(1,in_channels,1,1))
     
    
   

    

     

    
    

     

       self.beta = Parameter(torch.Tensor(1,in_channels,1,1))
     
    
   

    

     

    
    

     
 if nonlinear:
     
    
   

    

     

    
    

     

       self.v = Parameter(torch.Tensor(1,in_channels,1,1))
     
    
   

    

     

    
    

     

       self.reset_parameters()
     
    
   

    

     

    
    

     
 
     
    
   

    

     

    
    

     
 def reset_parameters(self):
     
    
   

    

     

    
    

     

       init.ones_(self.gamma)
     
    
   

    

     

    
    

     

       init.zeros_(self.beta)
     
    
   

    

     

    
    

     
 if self.nonlinear:
     
    
   

    

     

    
    

     

       init.ones_(self.v)
     
    
   

    

     

    
    

     
 def forward(self, x):
     
    
   

    

     

    
    

     
 if self.nonlinear:
     
    
   

    

     

    
    

     

       num = torch.sigmoid(self.v * x)
     
    
   

    

     

    
    

     

       std = group_std(x,self.groups)
     
    
   

    

     

    
    

     
 return num * std * self.gamma + self.beta
     
    
   

    

     

    
    

     
 else:
     
    
   

    

     

    
    

     
 return x * self.gamma + self.beta
     
    
  
 

文章来源: blog.csdn.net，作者：网奇，版权归原作者所有，如需转载，请联系作者。

原文链接：blog.csdn.net/jacke121/article/details/105434430