code: import torch.nn as nn
  import torch
  import torch.nn.functional as F
  from torch.distributions import Bernoulli
  
  
  def conv3x3(in_planes, out_planes, stride=1):
      """3x3 convolution with padding"""
      return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
                       padding=1, bias=False)
  
  
  class SELayer(nn.Module):
      def __init__(self, channel, reduction=16):
          super(SELayer, self).__init__()
          self.avg_pool = nn.AdaptiveAvgPool2d(1)
          self.fc = nn.Sequential(
              nn.Linear(channel, channel // reduction),
              nn.ReLU(inplace=True),
              nn.Linear(channel // reduction, channel),
              nn.Sigmoid()
          )
  
      def forward(self, x):
          b, c, _, _ = x.size()
          y = self.avg_pool(x).view(b, c)
          y = self.fc(y).view(b, c, 1, 1)
          return x * y
  
  
  class DropBlock(nn.Module):
      def __init__(self, block_size):
          super(DropBlock, self).__init__()
  
          self.block_size = block_size
          # self.gamma = gamma
          # self.bernouli = Bernoulli(gamma)
  
      def forward(self, x, gamma):
          device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
          # shape: (bsize, channels, height, width)
  
          if self.training:
              batch_size, channels, height, width = x.shape
  
              bernoulli = Bernoulli(gamma)
              mask = bernoulli.sample(
                  (batch_size, channels, height - (self.block_size - 1), width - (self.block_size - 1))).to(device)
              block_mask = self._compute_block_mask(mask)
              countM = block_mask.size()[0] * block_mask.size()[1] * block_mask.size()[2] * block_mask.size()[3]
              count_ones = block_mask.sum()
  
              return block_mask * x * (countM / count_ones)
          else:
              return x
  
      def _compute_block_mask(self, mask):
          device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
  
          left_padding = int((self.block_size - 1) / 2)
          right_padding = int(self.block_size / 2)
  
          batch_size, channels, height, width = mask.shape
          # print ("mask", mask[0][0])
          non_zero_idxs = mask.nonzero()
          nr_blocks = non_zero_idxs.shape[0]
  
          offsets = torch.stack(
              [
                  torch.arange(self.block_size).view(-1, 1).expand(self.block_size, self.block_size).reshape(-1),
                  # - left_padding,
                  torch.arange(self.block_size).repeat(self.block_size),  # - left_padding
              ]
          ).t().to(device)
          offsets = torch.cat((torch.zeros(self.block_size ** 2, 2).to(device).long(), offsets.long()), 1)
  
          if nr_blocks > 0:
              non_zero_idxs = non_zero_idxs.repeat(self.block_size ** 2, 1)
              offsets = offsets.repeat(nr_blocks, 1).view(-1, 4)
              offsets = offsets.long()
  
              block_idxs = non_zero_idxs + offsets
              # block_idxs += left_padding
              padded_mask = F.pad(mask, (left_padding, right_padding, left_padding, right_padding))
              padded_mask[block_idxs[:, 0], block_idxs[:, 1], block_idxs[:, 2], block_idxs[:, 3]] = 1.
          else:
              padded_mask = F.pad(mask, (left_padding, right_padding, left_padding, right_padding))
  
          block_mask = 1 - padded_mask  # [:height, :width]
          return block_mask
  
  
  class BasicBlock(nn.Module):
      expansion = 1
  
      def __init__(self, inplanes, planes, stride=1, downsample=None, drop_rate=0.0, drop_block=False,
                   block_size=1, use_se=False):
          super(BasicBlock, self).__init__()
          self.conv1 = conv3x3(inplanes, planes)
          self.bn1 = nn.BatchNorm2d(planes)
          self.relu = nn.LeakyReLU(0.1)
          self.conv2 = conv3x3(planes, planes)
          self.bn2 = nn.BatchNorm2d(planes)
          self.conv3 = conv3x3(planes, planes)
          self.bn3 = nn.BatchNorm2d(planes)
          self.maxpool = nn.MaxPool2d(stride)
          self.downsample = downsample
          self.stride = stride
          self.drop_rate = drop_rate
          self.num_batches_tracked = 0
          self.drop_block = drop_block
          self.block_size = block_size
          self.DropBlock = DropBlock(block_size=self.block_size)
          self.use_se = use_se
          if self.use_se:
              self.se = SELayer(planes, 4)
  
      def forward(self, x):
          self.num_batches_tracked += 1
  
          residual = x
  
          out = self.conv1(x)
          out = self.bn1(out)
          out = self.relu(out)
  
          out = self.conv2(out)
          out = self.bn2(out)
          out = self.relu(out)
  
          out = self.conv3(out)
          out = self.bn3(out)
          if self.use_se:
              out = self.se(out)
  
          if self.downsample is not None:
              residual = self.downsample(x)
          out += residual
          out = self.relu(out)
          out = self.maxpool(out)
  
          if self.drop_rate > 0:
              if self.drop_block == True:
                  feat_size = out.size()[2]
                  keep_rate = max(1.0 - self.drop_rate / (20 * 2000) * (self.num_batches_tracked), 1.0 - self.drop_rate)
                  gamma = (1 - keep_rate) / self.block_size ** 2 * feat_size ** 2 / (feat_size - self.block_size + 1) ** 2
                  out = self.DropBlock(out, gamma=gamma)
              else:
                  out = F.dropout(out, p=self.drop_rate, training=self.training, inplace=True)
  
          return out
  
  
  class ResNet(nn.Module):
  
      def __init__(self, block, n_blocks, keep_prob=1.0, avg_pool=False, drop_rate=0.0,
                   dropblock_size=5, num_classes=-1, use_se=False):
          super(ResNet, self).__init__()
  
          self.inplanes = 3
          self.use_se = use_se
          self.layer1 = self._make_layer(block, n_blocks[0], 64,
                                         stride=2, drop_rate=drop_rate)
          self.layer2 = self._make_layer(block, n_blocks[1], 160,
                                         stride=2, drop_rate=drop_rate)
          self.layer3 = self._make_layer(block, n_blocks[2], 320,
                                         stride=2, drop_rate=drop_rate, drop_block=True, block_size=dropblock_size)
          self.layer4 = self._make_layer(block, n_blocks[3], 640,
                                         stride=2, drop_rate=drop_rate, drop_block=True, block_size=dropblock_size)
          if avg_pool:
              # self.avgpool = nn.AvgPool2d(5, stride=1)
              self.avgpool = nn.AdaptiveAvgPool2d(1)
          self.keep_prob = keep_prob
          self.keep_avg_pool = avg_pool
          self.dropout = nn.Dropout(p=1 - self.keep_prob, inplace=False)
          self.drop_rate = drop_rate
  
          for m in self.modules():
              if isinstance(m, nn.Conv2d):
                  nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='leaky_relu')
              elif isinstance(m, nn.BatchNorm2d):
                  nn.init.constant_(m.weight, 1)
                  nn.init.constant_(m.bias, 0)
  
          self.num_classes = num_classes
          if self.num_classes > 0:
              self.classifier = nn.Linear(640, self.num_classes)
  
      def _make_layer(self, block, n_block, planes, stride=1, drop_rate=0.0, drop_block=False, block_size=1):
          downsample = None
          if stride != 1 or self.inplanes != planes * block.expansion:
              downsample = nn.Sequential(
                  nn.Conv2d(self.inplanes, planes * block.expansion,
                            kernel_size=1, stride=1, bias=False),
                  nn.BatchNorm2d(planes * block.expansion),
              )
  
          layers = []
          if n_block == 1:
              layer = block(self.inplanes, planes, stride, downsample, drop_rate, drop_block, block_size, self.use_se)
          else:
              layer = block(self.inplanes, planes, stride, downsample, drop_rate, self.use_se)
          layers.append(layer)
          self.inplanes = planes * block.expansion
  
          for i in range(1, n_block):
              if i == n_block - 1:
                  layer = block(self.inplanes, planes, drop_rate=drop_rate, drop_block=drop_block,
                                block_size=block_size, use_se=self.use_se)
              else:
                  layer = block(self.inplanes, planes, drop_rate=drop_rate, use_se=self.use_se)
              layers.append(layer)
  
          return nn.Sequential(*layers)
  
      def forward(self, x, is_feat=False):
          x = self.layer1(x)
          f0 = x
          x = self.layer2(x)
          f1 = x
          x = self.layer3(x)
          f2 = x
          x = self.layer4(x)
          f3 = x
          if self.keep_avg_pool:
              x = self.avgpool(x)
          x = x.view(x.size(0), -1)
          feat = x
          if self.num_classes > 0:
              x = self.classifier(x)
  
          if is_feat:
              return [f0, f1, f2, f3, feat], x
          else:
              return x
  
      def get_embedding(self, x):
          [f0, f1, f2, f3, feat], x = self.forward(x, is_feat=True)
          return feat
  
  
  def resnet12(keep_prob=1.0, avg_pool=False, **kwargs):
      """Constructs a ResNet-12 model.
      """
      model = ResNet(BasicBlock, [1, 1, 1, 1], keep_prob=keep_prob, avg_pool=avg_pool, **kwargs)
      return model
  
  
  def resnet18(keep_prob=1.0, avg_pool=False, **kwargs):
      """Constructs a ResNet-18 model.
      """
      model = ResNet(BasicBlock, [1, 1, 2, 2], keep_prob=keep_prob, avg_pool=avg_pool, **kwargs)
      return model
  
  
  def resnet24(keep_prob=1.0, avg_pool=False, **kwargs):
      """Constructs a ResNet-24 model.
      """
      model = ResNet(BasicBlock, [2, 2, 2, 2], keep_prob=keep_prob, avg_pool=avg_pool, **kwargs)
      return model
  
  
  def resnet50(keep_prob=1.0, avg_pool=False, **kwargs):
      """Constructs a ResNet-50 model.
      indeed, only (3 + 4 + 6 + 3) * 3 + 1 = 49 layers
      """
      model = ResNet(BasicBlock, [3, 4, 6, 3], keep_prob=keep_prob, avg_pool=avg_pool, **kwargs)
      return model
  
  
  def resnet101(keep_prob=1.0, avg_pool=False, **kwargs):
      """Constructs a ResNet-101 model.
      indeed, only (3 + 4 + 23 + 3) * 3 + 1 = 100 layers
      """
      model = ResNet(BasicBlock, [3, 4, 23, 3], keep_prob=keep_prob, avg_pool=avg_pool, **kwargs)
      return model
  
  
  def seresnet12(keep_prob=1.0, avg_pool=False, **kwargs):
      """Constructs a ResNet-12 model.
      """
      model = ResNet(BasicBlock, [1, 1, 1, 1], keep_prob=keep_prob, avg_pool=avg_pool, use_se=True, **kwargs)
      return model
  
  
  def seresnet18(keep_prob=1.0, avg_pool=False, **kwargs):
      """Constructs a ResNet-18 model.
      """
      model = ResNet(BasicBlock, [1, 1, 2, 2], keep_prob=keep_prob, avg_pool=avg_pool, use_se=True, **kwargs)
      return model
  
  
  def seresnet24(keep_prob=1.0, avg_pool=False, **kwargs):
      """Constructs a ResNet-24 model.
      """
      model = ResNet(BasicBlock, [2, 2, 2, 2], keep_prob=keep_prob, avg_pool=avg_pool, use_se=True, **kwargs)
      return model
  
  
  def seresnet50(keep_prob=1.0, avg_pool=False, **kwargs):
      """Constructs a ResNet-50 model.
      indeed, only (3 + 4 + 6 + 3) * 3 + 1 = 49 layers
      """
      model = ResNet(BasicBlock, [3, 4, 6, 3], keep_prob=keep_prob, avg_pool=avg_pool, use_se=True, **kwargs)
      return model
  
  
  def seresnet101(keep_prob=1.0, avg_pool=False, **kwargs):
      """Constructs a ResNet-101 model.
      indeed, only (3 + 4 + 23 + 3) * 3 + 1 = 100 layers
      """
      model = ResNet(BasicBlock, [3, 4, 23, 3], keep_prob=keep_prob, avg_pool=avg_pool, use_se=True, **kwargs)
      return model
  
  
  if __name__ == '__main__':
      from types import SimpleNamespace
      # import argparse
      #
      # parser = argparse.ArgumentParser('argument for training')
      # parser.add_argument('--model', type=str, choices=['resnet12', 'resnet18', 'resnet24', 'resnet50', 'resnet101',
      #                                                   'seresnet12', 'seresnet18', 'seresnet24', 'seresnet50',
      #                                                   'seresnet101'])
      # args = parser.parse_args()
      args = SimpleNamespace(model='resnet12')
      args.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
  
      model_dict = {
          'resnet12': resnet12,
          'resnet18': resnet18,
          'resnet24': resnet24,
          'resnet50': resnet50,
          'resnet101': resnet101,
          'seresnet12': seresnet12,
          'seresnet18': seresnet18,
          'seresnet24': seresnet24,
          'seresnet50': seresnet50,
          'seresnet101': seresnet101,
      }
  
      model = model_dict[args.model](avg_pool=True, drop_rate=0.1, dropblock_size=5, num_classes=64).to(args.device)
      data = torch.randn(2, 3, 84, 84)
      model = model.to(args.device)
      data = data.to(args.device)
      feat, logit = model(data, is_feat=True)
      print(feat[-1].shape)
      print(logit.shape)
  
      print("DONE")
  
apis: torch.cuda.is_available torch.nn.Conv2d torch.distributions.Bernoulli torch.nn.Sigmoid torch.nn.Dropout torch.nn.BatchNorm2d torch.nn.init.kaiming_normal_ torch.nn.functional.dropout torch.randn torch.nn.functional.pad torch.arange torch.nn.MaxPool2d torch.nn.AdaptiveAvgPool2d torch.nn.Linear torch.nn.init.constant_ torch.nn.Sequential torch.zeros torch.nn.ReLU torch.nn.LeakyReLU