lrn.py

import torch.nn as nn

class LRN(nn.Module):
  def __init__(self, size, alpha=1e-4, beta=0.75, k=1):
    super(LRN, self).__init__()
    self.avg = nn.AvgPool3d(kernel_size =(size,1,1), stride=1, padding=int((size-1)/2))
    self.alpha = alpha
    self.beta = beta
    self.k = k
    
  def forward(self, x):
    deno = x.pow(2).unsqueeze(1)
    deno = self.avg(deno).squeeze(1)
    deno = deno.mul(self.alpha).add(self.k).pow(self.beta)
    x = x.div(deno)
    return x 
  
  
x = torch.randn( 20,3, 10, 10)
lrn = nn.LocalResponseNorm(2)
lrn(x).size()