text: import torch
  import torch.nn as nn
  import torchvision
  from . import resnet, resnext, mobilenet, hrnet, shufflenetv2
  from lib.nn import SynchronizedBatchNorm2d
  BatchNorm2d = SynchronizedBatchNorm2d
  import torch.nn.functional as F
  from math import log2
  
  
  class SegmentationModuleBase(nn.Module):
      def __init__(self):
          super(SegmentationModuleBase, self).__init__()
  
      def pixel_acc(self, pred, label):
          _, preds = torch.max(pred, dim=1)
          valid = (label >= 0).long()
          acc_sum = torch.sum(valid * (preds == label).long())
          pixel_sum = torch.sum(valid)
          acc = acc_sum.float() / (pixel_sum.float() + 1e-10)
          return acc
  
  
  class SegmentationModule(SegmentationModuleBase):
      def __init__(self, net_enc, net_dec, crit, deep_sup_scale=None):
          super(SegmentationModule, self).__init__()
          self.encoder = net_enc
          self.decoder = net_dec
          self.crit = crit
          self.deep_sup_scale = deep_sup_scale
  
      def forward(self, feed_dict, *, segSize=None):
          # training
          
          if type(feed_dict) is list:
              feed_dict = feed_dict[0]
              # also, convert to torch.cuda.FloatTensor
              if torch.cuda.is_available():
                  feed_dict['img_data'] = feed_dict['img_data'].cuda()
                  feed_dict['seg_label'] = feed_dict['seg_label'].cuda()
              else:
                  raise RunTimeError('Cannot convert torch.Floattensor into torch.cuda.FloatTensor')
          if segSize is None:
              if self.deep_sup_scale is not None: # use deep supervision technique
                  (pred, pred_deepsup) = self.decoder(self.encoder(feed_dict['img_data'], return_feature_maps=True))
              else:
                  pred = self.decoder(self.encoder(feed_dict['img_data'], return_feature_maps=True))
  
              loss = self.crit(pred, feed_dict['seg_label'])
              if self.deep_sup_scale is not None:
                  loss_deepsup = self.crit(pred_deepsup, feed_dict['seg_label'])
                  loss = loss + loss_deepsup * self.deep_sup_scale
  
              acc = self.pixel_acc(pred, feed_dict['seg_label'])
              return loss, acc
          # inference
          else:
              pred = self.decoder(self.encoder(feed_dict['img_data'], return_feature_maps=True), segSize=segSize)
              return pred
  
  
  class ModelBuilder:
      # custom weights initialization
      @staticmethod
      def weights_init(m):
          classname = m.__class__.__name__
          if classname.find('Conv') != -1 and classname.find('_BNReluConv') == -1:
              nn.init.kaiming_normal_(m.weight.data)
          elif classname.find('BatchNorm') != -1:
              m.weight.data.fill_(1.)
              m.bias.data.fill_(1e-4)
          #elif classname.find('Linear') != -1:
          #    m.weight.data.normal_(0.0, 0.0001)
  
      @staticmethod
      def build_encoder(arch='resnet50dilated', fc_dim=512, weights=''):
          pretrained = True if len(weights) == 0 else False
          arch = arch.lower()
          if arch == 'mobilenetv2dilated':
              orig_mobilenet = mobilenet.__dict__['mobilenetv2'](pretrained=pretrained)
              net_encoder = MobileNetV2Dilated(orig_mobilenet, dilate_scale=8)
          elif arch == 'resnet18':
              orig_resnet = resnet.__dict__['resnet18'](pretrained=pretrained)
              net_encoder = Resnet(orig_resnet)
          elif arch == 'resnet18dilated':
              orig_resnet = resnet.__dict__['resnet18'](pretrained=pretrained)
              net_encoder = ResnetDilated(orig_resnet, dilate_scale=8)
          elif arch == 'resnet34':
              raise NotImplementedError
              orig_resnet = resnet.__dict__['resnet34'](pretrained=pretrained)
              net_encoder = Resnet(orig_resnet)
          elif arch == 'resnet34dilated':
              raise NotImplementedError
              orig_resnet = resnet.__dict__['resnet34'](pretrained=pretrained)
              net_encoder = ResnetDilated(orig_resnet, dilate_scale=8)
          elif arch == 'resnet50':
              orig_resnet = resnet.__dict__['resnet50'](pretrained=pretrained)
              net_encoder = Resnet(orig_resnet)
          elif arch == 'resnet50dilated':
              orig_resnet = resnet.__dict__['resnet50'](pretrained=pretrained)
              net_encoder = ResnetDilated(orig_resnet, dilate_scale=8)
          elif arch == 'resnet101':
              orig_resnet = resnet.__dict__['resnet101'](pretrained=pretrained)
              net_encoder = Resnet(orig_resnet)
          elif arch == 'resnet101dilated':
              orig_resnet = resnet.__dict__['resnet101'](pretrained=pretrained)
              net_encoder = ResnetDilated(orig_resnet, dilate_scale=8)
          elif arch == 'resnext101':
              orig_resnext = resnext.__dict__['resnext101'](pretrained=pretrained)
              net_encoder = Resnet(orig_resnext) # we can still use class Resnet
          elif arch == 'hrnetv2':
              net_encoder = hrnet.__dict__['hrnetv2'](pretrained=pretrained)
          else:
              net_encoder = shufflenetv2.__dict__['shufflenetv2'](pretrained=pretrained)
  
  
          # encoders are usually pretrained
          # net_encoder.apply(ModelBuilder.weights_init)
          if len(weights) > 0:
              print('Loading weights for net_encoder')
              net_encoder.load_state_dict(
                  torch.load(weights, map_location=lambda storage, loc: storage), strict=False)
          return net_encoder
  
      @staticmethod
      def build_decoder(arch='ppm_deepsup',
                        fc_dim=512, num_class=150,
                        weights='', use_softmax=False):
          arch = arch.lower()
          if arch == 'c1_deepsup':
              net_decoder = C1DeepSup(
                  num_class=num_class,
                  fc_dim=fc_dim,
                  use_softmax=use_softmax)
          elif arch == 'c1':
              net_decoder = C1(
                  num_class=num_class,
                  fc_dim=fc_dim,
                  use_softmax=use_softmax)
          elif arch == 'ppm':
              net_decoder = PPM(
                  num_class=num_class,
                  fc_dim=fc_dim,
                  use_softmax=use_softmax)
          elif arch == 'ppm_deepsup':
              net_decoder = PPMDeepsup(
                  num_class=num_class,
                  fc_dim=fc_dim,
                  use_softmax=use_softmax)
          elif arch == 'upernet_lite':
              net_decoder = UPerNet(
                  num_class=num_class,
                  fc_dim=fc_dim,
                  use_softmax=use_softmax,
                  fpn_dim=256)
          elif arch == 'upernet':
              net_decoder = UPerNet(
                  num_class=num_class,
                  fc_dim=fc_dim,
                  use_softmax=use_softmax,
                  fpn_dim=512)
          else:
              net_decoder = SwiftNet(
                  num_class=num_class,
                  use_softmax=use_softmax)
  
          net_decoder.apply(ModelBuilder.weights_init)
          if len(weights) > 0:
              print('Loading weights for net_decoder')
              net_decoder.load_state_dict(
                  torch.load(weights, map_location=lambda storage, loc: storage), strict=False)
          return net_decoder
  
  
  def conv3x3_bn_relu(in_planes, out_planes, stride=1):
      "3x3 convolution + BN + relu"
      return nn.Sequential(
              nn.Conv2d(in_planes, out_planes, kernel_size=3,
                        stride=stride, padding=1, bias=False),
              BatchNorm2d(out_planes),
              nn.ReLU(inplace=True),
              )
  
  
  class Resnet(nn.Module):
      def __init__(self, orig_resnet):
          super(Resnet, self).__init__()
  
          # take pretrained resnet, except AvgPool and FC
          self.conv1 = orig_resnet.conv1
          self.bn1 = orig_resnet.bn1
          self.relu1 = orig_resnet.relu1
          self.conv2 = orig_resnet.conv2
          self.bn2 = orig_resnet.bn2
          self.relu2 = orig_resnet.relu2
          self.conv3 = orig_resnet.conv3
          self.bn3 = orig_resnet.bn3
          self.relu3 = orig_resnet.relu3
          self.maxpool = orig_resnet.maxpool
          self.layer1 = orig_resnet.layer1
          self.layer2 = orig_resnet.layer2
          self.layer3 = orig_resnet.layer3
          self.layer4 = orig_resnet.layer4
  
      def forward(self, x, return_feature_maps=False):
          conv_out = []
  
          x = self.relu1(self.bn1(self.conv1(x)))
          x = self.relu2(self.bn2(self.conv2(x)))
          x = self.relu3(self.bn3(self.conv3(x)))
          x = self.maxpool(x)
  
          x = self.layer1(x); conv_out.append(x);
          x = self.layer2(x); conv_out.append(x);
          x = self.layer3(x); conv_out.append(x);
          x = self.layer4(x); conv_out.append(x);
  
          if return_feature_maps:
              return conv_out
          return [x]
  
  
  class ResnetDilated(nn.Module):
      def __init__(self, orig_resnet, dilate_scale=8):
          super(ResnetDilated, self).__init__()
          from functools import partial
  
          if dilate_scale == 8:
              orig_resnet.layer3.apply(
                  partial(self._nostride_dilate, dilate=2))
              orig_resnet.layer4.apply(
                  partial(self._nostride_dilate, dilate=4))
          elif dilate_scale == 16:
              orig_resnet.layer4.apply(
                  partial(self._nostride_dilate, dilate=2))
  
          # take pretrained resnet, except AvgPool and FC
          self.conv1 = orig_resnet.conv1
          self.bn1 = orig_resnet.bn1
          self.relu1 = orig_resnet.relu1
          self.conv2 = orig_resnet.conv2
          self.bn2 = orig_resnet.bn2
          self.relu2 = orig_resnet.relu2
          self.conv3 = orig_resnet.conv3
          self.bn3 = orig_resnet.bn3
          self.relu3 = orig_resnet.relu3
          self.maxpool = orig_resnet.maxpool
          self.layer1 = orig_resnet.layer1
          self.layer2 = orig_resnet.layer2
          self.layer3 = orig_resnet.layer3
          self.layer4 = orig_resnet.layer4
  
      def _nostride_dilate(self, m, dilate):
          classname = m.__class__.__name__
          if classname.find('Conv') != -1:
              # the convolution with stride
              if m.stride == (2, 2):
                  m.stride = (1, 1)
                  if m.kernel_size == (3, 3):
                      m.dilation = (dilate//2, dilate//2)
                      m.padding = (dilate//2, dilate//2)
              # other convoluions
              else:
                  if m.kernel_size == (3, 3):
                      m.dilation = (dilate, dilate)
                      m.padding = (dilate, dilate)
  
      def forward(self, x, return_feature_maps=False):
          conv_out = []
  
          x = self.relu1(self.bn1(self.conv1(x)))
          x = self.relu2(self.bn2(self.conv2(x)))
          x = self.relu3(self.bn3(self.conv3(x)))
          x = self.maxpool(x)
  
          x = self.layer1(x); conv_out.append(x);
          x = self.layer2(x); conv_out.append(x);
          x = self.layer3(x); conv_out.append(x);
          x = self.layer4(x); conv_out.append(x);
  
          if return_feature_maps:
              return conv_out
          return [x]
  
  
  class MobileNetV2Dilated(nn.Module):
      def __init__(self, orig_net, dilate_scale=8):
          super(MobileNetV2Dilated, self).__init__()
          from functools import partial
  
          # take pretrained mobilenet features
          self.features = orig_net.features[:-1]
  
          self.total_idx = len(self.features)
          self.down_idx = [2, 4, 7, 14]
  
          if dilate_scale == 8:
              for i in range(self.down_idx[-2], self.down_idx[-1]):
                  self.features[i].apply(
                      partial(self._nostride_dilate, dilate=2)
                  )
              for i in range(self.down_idx[-1], self.total_idx):
                  self.features[i].apply(
                      partial(self._nostride_dilate, dilate=4)
                  )
          elif dilate_scale == 16:
              for i in range(self.down_idx[-1], self.total_idx):
                  self.features[i].apply(
                      partial(self._nostride_dilate, dilate=2)
                  )
  
      def _nostride_dilate(self, m, dilate):
          classname = m.__class__.__name__
          if classname.find('Conv') != -1:
              # the convolution with stride
              if m.stride == (2, 2):
                  m.stride = (1, 1)
                  if m.kernel_size == (3, 3):
                      m.dilation = (dilate//2, dilate//2)
                      m.padding = (dilate//2, dilate//2)
              # other convoluions
              else:
                  if m.kernel_size == (3, 3):
                      m.dilation = (dilate, dilate)
                      m.padding = (dilate, dilate)
  
      def forward(self, x, return_feature_maps=False):
          if return_feature_maps:
              conv_out = []
              for i in range(self.total_idx):
                  x = self.features[i](x)
                  if i in self.down_idx:
                      conv_out.append(x)
              conv_out.append(x)
              return conv_out
  
          else:
              return [self.features(x)]
  
  
  # last conv, deep supervision
  class C1DeepSup(nn.Module):
      def __init__(self, num_class=150, fc_dim=2048, use_softmax=False):
          super(C1DeepSup, self).__init__()
          self.use_softmax = use_softmax
  
          self.cbr = conv3x3_bn_relu(fc_dim, fc_dim // 4, 1)
          self.cbr_deepsup = conv3x3_bn_relu(fc_dim // 2, fc_dim // 4, 1)
  
          # last conv
          self.conv_last = nn.Conv2d(fc_dim // 4, num_class, 1, 1, 0)
          self.conv_last_deepsup = nn.Conv2d(fc_dim // 4, num_class, 1, 1, 0)
  
      def forward(self, conv_out, segSize=None):
          conv5 = conv_out[-1]
  
          x = self.cbr(conv5)
          x = self.conv_last(x)
  
          if self.use_softmax:  # is True during inference
              x = nn.functional.interpolate(
                  x, size=segSize, mode='bilinear', align_corners=False)
              x = nn.functional.softmax(x, dim=1)
              return x
  
          # deep sup
          conv4 = conv_out[-2]
          _ = self.cbr_deepsup(conv4)
          _ = self.conv_last_deepsup(_)
  
          x = nn.functional.log_softmax(x, dim=1)
          _ = nn.functional.log_softmax(_, dim=1)
  
          return (x, _)
  
  
  # last conv
  class C1(nn.Module):
      def __init__(self, num_class=150, fc_dim=2048, use_softmax=False):
          super(C1, self).__init__()
          self.use_softmax = use_softmax
  
          self.cbr = conv3x3_bn_relu(fc_dim, fc_dim // 4, 1)
  
          # last conv
          self.conv_last = nn.Conv2d(fc_dim // 4, num_class, 1, 1, 0)
  
      def forward(self, conv_out, segSize=None):
          conv5 = conv_out[-1]
          x = self.cbr(conv5)
          x = self.conv_last(x)
          if self.use_softmax: # is True during inference
              x = nn.functional.interpolate(
                  x, size=segSize, mode='bilinear', align_corners=False)
              x = nn.functional.softmax(x, dim=1)
  
          return x
  
  
  # pyramid pooling
  class PPM(nn.Module):
      def __init__(self, num_class=150, fc_dim=4096,
                   use_softmax=False, pool_scales=(1, 2, 3, 6)):
          super(PPM, self).__init__()
          self.use_softmax = use_softmax
  
          self.ppm = []
          for scale in pool_scales:
              self.ppm.append(nn.Sequential(
                  nn.AdaptiveAvgPool2d(scale),
                  nn.Conv2d(fc_dim, 512, kernel_size=1, bias=False),
                  BatchNorm2d(512),
                  nn.ReLU(inplace=True)
              ))
          self.ppm = nn.ModuleList(self.ppm)
  
          self.conv_last = nn.Sequential(
              nn.Conv2d(fc_dim+len(pool_scales)*512, 512,
                        kernel_size=3, padding=1, bias=False),
              BatchNorm2d(512),
              nn.ReLU(inplace=True),
              nn.Dropout2d(0.1),
              nn.Conv2d(512, num_class, kernel_size=1)
          )
  
      def forward(self, conv_out, segSize=None):
          conv5 = conv_out[-1]
  
          input_size = conv5.size()
          ppm_out = [conv5]
          for pool_scale in self.ppm:
              ppm_out.append(nn.functional.interpolate(
                  pool_scale(conv5),
                  (input_size[2], input_size[3]),
                  mode='bilinear', align_corners=False))
          ppm_out = torch.cat(ppm_out, 1)
  
          x = self.conv_last(ppm_out)
  
          if self.use_softmax:  # is True during inference
              x = nn.functional.interpolate(
                  x, size=segSize, mode='bilinear', align_corners=False)
              x = nn.functional.softmax(x, dim=1)
          else:
              x = nn.functional.log_softmax(x, dim=1)
          return x
  
  
  # pyramid pooling, deep supervision
  class PPMDeepsup(nn.Module):
      def __init__(self, num_class=150, fc_dim=4096,
                   use_softmax=False, pool_scales=(1, 2, 3, 6)):
          super(PPMDeepsup, self).__init__()
          self.use_softmax = use_softmax
  
          self.ppm = []
          for scale in pool_scales:
              self.ppm.append(nn.Sequential(
                  nn.AdaptiveAvgPool2d(scale),
                  nn.Conv2d(fc_dim, 512, kernel_size=1, bias=False),
                  BatchNorm2d(512),
                  nn.ReLU(inplace=True)
              ))
          self.ppm = nn.ModuleList(self.ppm)
          self.cbr_deepsup = conv3x3_bn_relu(fc_dim // 2, fc_dim // 4, 1)
  
          self.conv_last = nn.Sequential(
              nn.Conv2d(fc_dim+len(pool_scales)*512, 512,
                        kernel_size=3, padding=1, bias=False),
              BatchNorm2d(512),
              nn.ReLU(inplace=True),
              nn.Dropout2d(0.1),
              nn.Conv2d(512, num_class, kernel_size=1)
          )
          self.conv_last_deepsup = nn.Conv2d(fc_dim // 4, num_class, 1, 1, 0)
          self.dropout_deepsup = nn.Dropout2d(0.1)
  
      def forward(self, conv_out, segSize=None):
          conv5 = conv_out[-1]
  
          input_size = conv5.size()
          ppm_out = [conv5]
          for pool_scale in self.ppm:
              ppm_out.append(nn.functional.interpolate(
                  pool_scale(conv5),
                  (input_size[2], input_size[3]),
                  mode='bilinear', align_corners=False))
          ppm_out = torch.cat(ppm_out, 1)
  
          x = self.conv_last(ppm_out)
  
          if self.use_softmax:  # is True during inference
              x = nn.functional.interpolate(
                  x, size=segSize, mode='bilinear', align_corners=False)
              x = nn.functional.softmax(x, dim=1)
              return x
  
          # deep sup
          conv4 = conv_out[-2]
          _ = self.cbr_deepsup(conv4)
          _ = self.dropout_deepsup(_)
          _ = self.conv_last_deepsup(_)
  
          x = nn.functional.log_softmax(x, dim=1)
          _ = nn.functional.log_softmax(_, dim=1)
  
          return (x, _)
  
  
  # upernet
  class UPerNet(nn.Module):
      def __init__(self, num_class=6, fc_dim=4096,
                   use_softmax=False, pool_scales=(1, 2, 3, 6),
                   fpn_inplanes=(232, 464), fpn_dim=64):
          super(UPerNet, self).__init__()
          self.use_softmax = use_softmax
  
          # PPM Module
          self.ppm_pooling = []
          self.ppm_conv = []
  
          for scale in pool_scales:
              self.ppm_pooling.append(nn.AdaptiveAvgPool2d(scale))
              self.ppm_conv.append(nn.Sequential(
                  nn.Conv2d(fc_dim, 512, kernel_size=1, bias=False),
                  BatchNorm2d(512),
                  nn.ReLU(inplace=True)
              ))
          self.ppm_pooling = nn.ModuleList(self.ppm_pooling)
          self.ppm_conv = nn.ModuleList(self.ppm_conv)
          self.ppm_last_conv = conv3x3_bn_relu(fc_dim + len(pool_scales)*512, fpn_dim, 1)
  
          # FPN Module
          self.fpn_in = []
          for fpn_inplane in fpn_inplanes[:-1]:   # skip the top layer
              self.fpn_in.append(nn.Sequential(
                  nn.Conv2d(fpn_inplane, fpn_dim, kernel_size=1, bias=False),
                  BatchNorm2d(fpn_dim),
                  nn.ReLU(inplace=True)
              ))
          self.fpn_in = nn.ModuleList(self.fpn_in)
  
          self.fpn_out = []
          for i in range(len(fpn_inplanes) - 1):  # skip the top layer
              self.fpn_out.append(nn.Sequential(
                  conv3x3_bn_relu(fpn_dim, fpn_dim, 1),
              ))
          self.fpn_out = nn.ModuleList(self.fpn_out)
          self.conv_last = nn.Sequential(
              conv3x3_bn_relu(512, fpn_dim, 1),
              nn.Conv2d(fpn_dim, num_class, kernel_size=1)
          )
  
      def forward(self, conv_out, segSize=None):
          conv5 = conv_out[-1]
  
          input_size = conv5.size()
          ppm_out = [conv5]
          for pool_scale, pool_conv in zip(self.ppm_pooling, self.ppm_conv):
              ppm_out.append(pool_conv(nn.functional.interpolate(
                  pool_scale(conv5),
                  (input_size[2], input_size[3]),
                  mode='bilinear', align_corners=False)))
          ppm_out = torch.cat(ppm_out, 1)
          f = self.ppm_last_conv(ppm_out)
  
          fpn_feature_list = [f]
          for i in reversed(range(len(conv_out) - 1)):
              conv_x = conv_out[i]
              conv_x = self.fpn_in[i](conv_x) # lateral branch
  
              f = nn.functional.interpolate(
                  f, size=conv_x.size()[2:], mode='bilinear', align_corners=False) # top-down branch
              f = conv_x + f
  
              fpn_feature_list.append(self.fpn_out[i](f))
  
          fpn_feature_list.reverse() # [P2 - P5]
          output_size = fpn_feature_list[0].size()[2:]
          fusion_list = [fpn_feature_list[0]]
          for i in range(1, len(fpn_feature_list)):
              fusion_list.append(nn.functional.interpolate(
                  fpn_feature_list[i],
                  output_size,
                  mode='bilinear', align_corners=False))
          fusion_out = torch.cat(fusion_list, 1)
          x = self.conv_last(fusion_out)
  
          if self.use_softmax:  # is True during inference
              x = nn.functional.interpolate(
                  x, size=segSize, mode='bilinear', align_corners=False)
              x = nn.functional.softmax(x, dim=1)
              return x
  
          x = nn.functional.log_softmax(x, dim=1)
  
          return x
  
  upsample = lambda x, size: F.interpolate(x, size, mode='bilinear', align_corners=False)
  
  
  class SeparableConv2d(nn.Module):
      def __init__(self, in_channels, out_channels, kernel_size=1, stride=1, padding=0, dilation=1, bias=False):
          super(SeparableConv2d, self).__init__()
  
          self.conv1 = nn.Conv2d(in_channels, in_channels, kernel_size, stride, padding, dilation, groups=in_channels,
                                 bias=bias)
          self.pointwise = nn.Conv2d(in_channels, out_channels, 1, 1, 0, 1, 1, bias=bias)
  
      def forward(self, x):
          x = self.conv1(x)
          x = self.pointwise(x)
          return x
  
  class _BNReluConv(nn.Sequential):
      def __init__(self, num_maps_in, num_maps_out, k=3, batch_norm=True, bn_momentum=0.1, bias=False, dilation=1):
          super(_BNReluConv, self).__init__()
          if batch_norm:
              self.add_module('norm', nn.BatchNorm2d(num_maps_in, momentum=bn_momentum))
          self.add_module('relu', nn.ReLU(inplace=batch_norm is True))
          padding = k // 2
          self.add_module('conv', nn.Conv2d(num_maps_in, num_maps_out,
                                            kernel_size=k, padding=padding, bias=bias, dilation=dilation))
  
  class _Upsample(nn.Module):
      def __init__(self, num_maps_in, skip_maps_in, num_maps_out, use_bn=True, k=3):
          super(_Upsample, self).__init__()
          print(f'Upsample layer: in = {num_maps_in}, skip = {skip_maps_in}, out = {num_maps_out}')
          self.bottleneck = _BNReluConv(skip_maps_in, num_maps_in, k=1, batch_norm=use_bn)
          self.blend_conv = _BNReluConv(num_maps_in, num_maps_out, k=k, batch_norm=use_bn)
  
      def forward(self, x, skip):
          skip = self.bottleneck.forward(skip)
          skip_size = skip.size()[2:4]
          x = upsample(x, skip_size)
          x = x + skip
          x = self.blend_conv.forward(x)
          return x
  
  class SpatialPyramidPooling(nn.Module):
      def __init__(self, num_maps_in, num_levels, bt_size=512, level_size=128, out_size=128,
                   grids=(6, 3, 2, 1), square_grid=False, bn_momentum=0.1):
          super(SpatialPyramidPooling, self).__init__()
          self.grids = grids
          self.square_grid = square_grid
          self.spp = nn.Sequential()
          self.spp.add_module('spp_bn',
                              _BNReluConv(num_maps_in, bt_size, k=1, bn_momentum=bn_momentum, batch_norm=True))
          num_features = bt_size
          final_size = num_features
          for i in range(num_levels):
              final_size += level_size
              self.spp.add_module('spp' + str(i),
                                  _BNReluConv(num_features, level_size, k=1, bn_momentum=bn_momentum, batch_norm=True))
          self.spp.add_module('spp_fuse',
                              _BNReluConv(final_size, out_size, k=1, bn_momentum=bn_momentum, batch_norm=True))
  
      def forward(self, x):
          levels = []
          target_size = x.size()[2:4]
  
          ar = target_size[1] / target_size[0]
  
          x = self.spp[0].forward(x)
          levels.append(x)
          num = len(self.spp) - 1
  
          for i in range(1, num):
              if not self.square_grid:
                  grid_size = (self.grids[i - 1], max(1, round(ar * self.grids[i - 1])))
                  x_pooled = F.adaptive_avg_pool2d(x, grid_size)
              else:
                  x_pooled = F.adaptive_avg_pool2d(x, self.grids[i - 1])
              level = self.spp[i].forward(x_pooled)
  
              level = upsample(level, target_size)
              levels.append(level)
          x = torch.cat(levels, 1)
          x = self.spp[-1].forward(x)
          return x
  
  
  class SwiftNet(nn.Module):
      def __init__(self, num_class=150, use_softmax=False):
          super(SwiftNet, self).__init__()
          self.use_softmax = use_softmax
  
          upsamples = []
          upsamples.append(_Upsample(128, 116, 128, k=3))
          upsamples.append(_Upsample(128, 232, 128, k=3))
  
          self.upsample = nn.ModuleList(list(reversed(upsamples)))
          self.spp = SpatialPyramidPooling(464, 3, bt_size=6, level_size=42,
                  out_size=128, grids=(8, 4, 2, 1),
                  bn_momentum=0.01 / 2)
          self.logits = _BNReluConv(128, 6, batch_norm=True, k=1, bias=True   )
      def forward(self, conv_out, segSize=None):
          last_conv = conv_out.pop()
          conv_out.append(self.spp.forward(last_conv))
          conv_out = conv_out[::-1]
  
          x = conv_out[0]
  
          for skip, up in zip(conv_out[1:], self.upsample):
              x = up(x, skip)
  
          x = self.logits(x)
          
          if self.use_softmax: # is True during inference
              x = nn.functional.interpolate(
                  x, size=segSize, mode='bilinear', align_corners=False)
              x = nn.functional.softmax(x, dim=1)
  
          return x