text: import torch
  from torch import nn
  from torch.nn import functional as F
  import math
  
  
  class LayerNorm(nn.Module):
      def __init__(self, eps=1e-6):
          super().__init__()
          self.eps = eps
  
      def forward(self, x):
          mean = x.mean(-1, keepdim=True)
          std = x.std(-1, keepdim=True)
          return (x - mean) / (std + self.eps)
  
  
  class _NonLocalBlockND(nn.Module):
      def __init__(
          self,
          in_channels,
          inter_channels=None,
          dimension=3,
          sub_sample=True,
          bn_layer=True,
          dropout=0.2,
      ):
          super(_NonLocalBlockND, self).__init__()
  
          assert dimension in [1, 2, 3]
  
          self.dimension = dimension
          self.sub_sample = sub_sample
  
          self.in_channels = in_channels
          self.inter_channels = inter_channels
  
          if self.inter_channels is None:
              self.inter_channels = in_channels // 2
              if self.inter_channels == 0:
                  self.inter_channels = 1
  
          if dimension == 3:
              conv_nd = nn.Conv3d
              max_pool_layer = nn.MaxPool3d(kernel_size=(1, 2, 2))
              bn = nn.BatchNorm3d
          elif dimension == 2:
              conv_nd = nn.Conv2d
              max_pool_layer = nn.MaxPool2d(kernel_size=(2, 2))
              bn = nn.BatchNorm2d
          else:
              conv_nd = nn.Conv1d
              max_pool_layer = nn.MaxPool1d(kernel_size=(2))
              bn = nn.BatchNorm1d
  
          self.g = conv_nd(
              in_channels=self.in_channels,
              out_channels=self.inter_channels,
              kernel_size=1,
              stride=1,
              padding=0,
          )
  
          if bn_layer:
              self.W = nn.Sequential(
                  conv_nd(
                      in_channels=self.inter_channels,
                      out_channels=self.in_channels,
                      kernel_size=1,
                      stride=1,
                      padding=0,
                  ),
                  bn(self.in_channels),
              )
              nn.init.constant_(self.W[1].weight, 0)
              nn.init.constant_(self.W[1].bias, 0)
          else:
              self.W = conv_nd(
                  in_channels=self.inter_channels,
                  out_channels=self.in_channels,
                  kernel_size=1,
                  stride=1,
                  padding=0,
              )
              nn.init.constant_(self.W.weight, 0)
              nn.init.constant_(self.W.bias, 0)
  
          self.theta = conv_nd(
              in_channels=self.in_channels,
              out_channels=self.inter_channels,
              kernel_size=1,
              stride=1,
              padding=0,
          )
          self.phi = conv_nd(
              in_channels=self.in_channels,
              out_channels=self.inter_channels,
              kernel_size=1,
              stride=1,
              padding=0,
          )
  
          self.dropout = nn.Dropout(dropout)
          self.ln = LayerNorm()
  
          if sub_sample:
              self.g = nn.Sequential(self.g, max_pool_layer)
              self.phi = nn.Sequential(self.phi, max_pool_layer)
  
      def forward(self, x_support, query):
          """
          :param x_support: (b, c, t, h, w)
          :return:
          """
  
          x_support = x_support.permute(1, 0).unsqueeze(0)
          query = query.permute(1, 0).unsqueeze(0)
  
          batch_size = x_support.size(0)
  
          g_x = self.g(query).view(batch_size, self.inter_channels, -1)
          g_x = g_x.permute(0, 2, 1)  # BxHWxC
  
          theta_x = self.theta(x_support).view(
              batch_size, self.inter_channels, -1
          )  # BxCxHW
          theta_x = theta_x.permute(0, 2, 1)  # BxHWxC
          phi_x = self.phi(query).view(
              batch_size, self.inter_channels, -1
          )  # BxCxHW
          f = torch.matmul(theta_x, phi_x)  # BxHWxHW
          f = f / math.sqrt(self.inter_channels)  # rescale
          f_div_C = F.softmax(f, dim=-1)
  
          y = torch.matmul(f_div_C, g_x)  # BxHWxC
          y = self.ln(y)  # layer normalization in last dim
          y = y.permute(0, 2, 1).contiguous()  # BxCxHW
          y = y.view(batch_size, self.inter_channels, *x_support.size()[2:])
          W_y = self.dropout(self.W(F.relu(y)))
  
          z = W_y + x_support
  
          z = z.squeeze(0).permute(1, 0)
  
          return z
  
  
  class NONLocalBlock1D(_NonLocalBlockND):
      def __init__(
          self, in_channels, inter_channels=None, sub_sample=True, bn_layer=True
      ):
          super(NONLocalBlock1D, self).__init__(
              in_channels,
              inter_channels=inter_channels,
              dimension=1,
              sub_sample=sub_sample,
              bn_layer=bn_layer,
          )
  
  
  class NONLocalBlock2D(_NonLocalBlockND):
      def __init__(
          self, in_channels, inter_channels=None, sub_sample=True, bn_layer=True
      ):
          super(NONLocalBlock2D, self).__init__(
              in_channels,
              inter_channels=inter_channels,
              dimension=2,
              sub_sample=sub_sample,
              bn_layer=bn_layer,
          )
  
  
  class NONLocalBlock3D(_NonLocalBlockND):
      def __init__(
          self, in_channels, inter_channels=None, sub_sample=True, bn_layer=True
      ):
          super(NONLocalBlock3D, self).__init__(
              in_channels,
              inter_channels=inter_channels,
              dimension=3,
              sub_sample=sub_sample,
              bn_layer=bn_layer,
          )
  
  
  if __name__ == "__main__":
      from torch.autograd import Variable
      import torch
  
      sub_sample = True
      bn_layer = True
  
      img = Variable(torch.zeros(2, 3, 20))
      net = NONLocalBlock1D(3, sub_sample=sub_sample, bn_layer=bn_layer)
      out = net(img)
      print(out.size())
  
      img = Variable(torch.zeros(2, 3, 20, 20))
      net = NONLocalBlock2D(3, sub_sample=sub_sample, bn_layer=bn_layer)
      out = net(img)
      print(out.size())
  
      img = Variable(torch.randn(2, 3, 10, 20, 20))
      net = NONLocalBlock3D(3, sub_sample=sub_sample, bn_layer=bn_layer)
      out = net(img)
      print(out.size())
  