text: from collections import OrderedDict
  
  import torch
  from torch import nn
  from torch.nn import functional as F
  
  
  class _SimpleSegmentationModel(nn.Module):
      def __init__(self, backbone, classifier, aux_classifier=None):
          super(_SimpleSegmentationModel, self).__init__()
          self.backbone = backbone
          self.classifier = classifier
          self.aux_classifier = aux_classifier
  
      def forward(self, x):
          input_shape = x.shape[-2:]
          # contract: features is a dict of tensors
          features = self.backbone(x)
  
          result = OrderedDict()
          x = features["out"]
          x = self.classifier(x)
          x = F.interpolate(x, size=input_shape, mode='bilinear', align_corners=False)
          result["out"] = x
  
          if self.aux_classifier is not None:
              x = features["aux"]
              x = self.aux_classifier(x)
              x = F.interpolate(x, size=input_shape, mode='bilinear', align_corners=False)
              result["aux"] = x
  
          return result
  
  
  class FCN(_SimpleSegmentationModel):
      pass
  
  
  class DeepLabV3(_SimpleSegmentationModel):
      pass
  
  
  class FCNHead(nn.Sequential):
      def __init__(self, in_channels, channels):
          inter_channels = in_channels // 4
          layers = [
              nn.Conv2d(in_channels, inter_channels, 3, padding=1, bias=False),
              nn.BatchNorm2d(inter_channels),
              nn.ReLU(),
              nn.Dropout(0.1),
              nn.Conv2d(inter_channels, channels, 1)
          ]
  
          super(FCNHead, self).__init__(*layers)
          """
          for m in self.modules():
              if isinstance(m, nn.Conv2d):
                  nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
              elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)):
                  nn.init.constant_(m.weight, 1)
                  nn.init.constant_(m.bias, 0)
          """
  
  
  class DeepLabHead(nn.Sequential):
      def __init__(self, in_channels, num_classes):
          super(DeepLabHead, self).__init__(
              ASPP(in_channels, [12, 24, 36]),
              nn.Conv2d(256, 256, 3, padding=1, bias=False),
              nn.BatchNorm2d(256),
              nn.ReLU(),
              nn.Conv2d(256, num_classes, 1)
          )
          """
          for m in self.modules():
              if isinstance(m, nn.Conv2d):
                  nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
              elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)):
                  nn.init.constant_(m.weight, 1)
                  nn.init.constant_(m.bias, 0)
          """
  
  
  class ASPPConv(nn.Sequential):
      def __init__(self, in_channels, out_channels, dilation):
          modules = [
              nn.Conv2d(in_channels, out_channels, 3, padding=dilation, dilation=dilation, bias=False),
              nn.BatchNorm2d(out_channels),
              nn.ReLU()
          ]
          super(ASPPConv, self).__init__(*modules)
  
  
  class ASPPPooling(nn.Sequential):
      def __init__(self, in_channels, out_channels):
          super(ASPPPooling, self).__init__(
              nn.AdaptiveAvgPool2d(1),
              nn.Conv2d(in_channels, out_channels, 1, bias=False),
              nn.BatchNorm2d(out_channels),
              nn.ReLU())
  
      def forward(self, x):
          size = x.shape[-2:]
          x = super(ASPPPooling, self).forward(x)
          return F.interpolate(x, size=size, mode='bilinear', align_corners=False)
  
  
  class ASPP(nn.Module):
      def __init__(self, in_channels, atrous_rates):
          super(ASPP, self).__init__()
          out_channels = 256
          modules = []
          modules.append(nn.Sequential(
              nn.Conv2d(in_channels, out_channels, 1, bias=False),
              nn.BatchNorm2d(out_channels),
              nn.ReLU()))
  
          rate1, rate2, rate3 = tuple(atrous_rates)
          modules.append(ASPPConv(in_channels, out_channels, rate1))
          modules.append(ASPPConv(in_channels, out_channels, rate2))
          modules.append(ASPPConv(in_channels, out_channels, rate3))
          modules.append(ASPPPooling(in_channels, out_channels))
  
          self.convs = nn.ModuleList(modules)
  
          self.project = nn.Sequential(
              nn.Conv2d(5 * out_channels, out_channels, 1, bias=False),
              nn.BatchNorm2d(out_channels),
              nn.ReLU(),
              nn.Dropout(0.5))
  
      def forward(self, x):
          res = []
          for conv in self.convs:
              res.append(conv(x))
          res = torch.cat(res, dim=1)
          return self.project(res)
  