text: # ------------------------------------------------------------------------------
  # Reference: https://github.com/pytorch/vision/blob/master/torchvision/models/segmentation/deeplabv3.py
  # Modified by Bowen Cheng (bcheng9@illinois.edu)
  # ------------------------------------------------------------------------------
  
  import torch
  from torch import nn
  from torch.nn import functional as F
  
  from .conv_module import depthwise_separable_conv
  
  __all__ = ["ASPP", "SepASPP"]
  
  
  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(inplace=True)
          ]
          super(ASPPConv, self).__init__(*modules)
  
  
  class SepASPPConv(nn.Sequential):
      def __init__(self, in_channels, out_channels, dilation):
          modules = depthwise_separable_conv(in_channels, out_channels, 3, padding=dilation, dilation=dilation)
          super(SepASPPConv, self).__init__(*modules)
  
  
  class ASPPPooling(nn.Module):
      def __init__(self, in_channels, out_channels):
          super(ASPPPooling, self).__init__()
          self.aspp_pooling = nn.Sequential(
              nn.AdaptiveAvgPool2d(1),
              nn.Conv2d(in_channels, out_channels, 1, bias=False),
              nn.ReLU(inplace=True)
          )
  
      def set_image_pooling(self, pool_size):
          self.aspp_pooling[0] = nn.AvgPool2d(kernel_size=pool_size, stride=1)
  
      def forward(self, x):
          size = x.shape[-2:]
          x = self.aspp_pooling(x)
          return F.interpolate(x, size=size, mode='bilinear', align_corners=True)
  
  
  class ASPP(nn.Module):
      def __init__(self, in_channels, out_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(inplace=True)))
  
          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(inplace=True),
              nn.Dropout(0.1))
  
      def set_image_pooling(self, pool_size):
          self.convs[-1].set_image_pooling(pool_size)
  
      def forward(self, x):
          res = []
          for conv in self.convs:
              res.append(conv(x))
          res = torch.cat(res, dim=1)
          return self.project(res)
  
  
  class SepASPP(nn.Module):
      def __init__(self, in_channels, out_channels, atrous_rates):
          super(SepASPP, 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(inplace=True)))
  
          rate1, rate2, rate3 = tuple(atrous_rates)
          modules.append(SepASPPConv(in_channels, out_channels, rate1))
          modules.append(SepASPPConv(in_channels, out_channels, rate2))
          modules.append(SepASPPConv(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(inplace=True),
              nn.Dropout(0.1))
  
      def set_image_pooling(self, pool_size):
          self.convs[-1].set_image_pooling(pool_size)
  
      def forward(self, x):
          res = []
          for conv in self.convs:
              res.append(conv(x))
          res = torch.cat(res, dim=1)
          return self.project(res)
  