text: # BSD 3-Clause License
  #
  # Copyright (c) Soumith Chintala 2016,
  # All rights reserved.
  #
  # Redistribution and use in source and binary forms, with or without
  # modification, are permitted provided that the following conditions are met:
  #
  # * Redistributions of source code must retain the above copyright notice, this
  #   list of conditions and the following disclaimer.
  #
  # * Redistributions in binary form must reproduce the above copyright notice,
  #   this list of conditions and the following disclaimer in the documentation
  #   and/or other materials provided with the distribution.
  #
  # * Neither the name of the copyright holder nor the names of its
  #   contributors may be used to endorse or promote products derived from
  #   this software without specific prior written permission.
  
  """Code for ResNet.
  Adapted from
  https://github.com/pytorch/vision/blob/master/torchvision/models/resnet.py
  """
  
  import math
  
  import torch.nn as nn
  import torch.utils.model_zoo as model_zoo
  from torchvision.models.resnet import BasicBlock
  from torchvision.models.vgg import cfg, model_urls, VGG
  
  
  class ResNet(nn.Module):
  
    def __init__(self, block, layers, in_channels, n_classes):
      self.inplanes = 64
      super(ResNet, self).__init__()
      self.conv1 = nn.Conv2d(
          in_channels, 64, kernel_size=7, stride=2, padding=3, bias=False)
      self.bn1 = nn.BatchNorm2d(64)
      self.relu = nn.ReLU(inplace=True)
      self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
      self.layer1 = self._make_layer(block, 64, layers[0])
      self.layer2 = self._make_layer(block, 128, layers[1], stride=2)
      self.layer3 = self._make_layer(block, 256, layers[2], stride=2)
      self.layer4 = self._make_layer(block, 512, layers[3], stride=2)
      self.avgpool = nn.AvgPool2d(7)
      self.fc = nn.Linear(512 * block.expansion, n_classes)
  
      for m in self.modules():
        if isinstance(m, nn.Conv2d):
          n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
          m.weight.data.normal_(0, math.sqrt(2. / n))
        elif isinstance(m, nn.BatchNorm2d):
          m.weight.data.fill_(1)
          m.bias.data.zero_()
  
    def _make_layer(self, block, planes, blocks, stride=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=stride,
                bias=False),
            nn.BatchNorm2d(planes * block.expansion),
        )
  
      layers = []
      layers.append(block(self.inplanes, planes, stride, downsample))
      self.inplanes = planes * block.expansion
      for i in range(1, blocks):
        layers.append(block(self.inplanes, planes))
  
      return nn.Sequential(*layers)
  
    def forward(self, x):
      x = self.conv1(x)
      x = self.bn1(x)
      x = self.relu(x)
      x = self.maxpool(x)
  
      x = self.layer1(x)
      x = self.layer2(x)
      x = self.layer3(x)
      conv4 = self.layer4(x)
  
      repr = self.avgpool(conv4).view(conv4.size(0), -1)
      logit = self.fc(repr)
  
      return logit, repr
  
  
  def resnet18(**kwargs):
    """Constructs a ResNet-18 model.
    Args:
        pretrained (bool): If True, returns a model pre-trained on ImageNet
    """
    model = ResNet(BasicBlock, [2, 2, 2, 2], **kwargs)
    return model
  
  
  def get_resnet(in_channels, n_classes):
    return resnet18(in_channels=in_channels, n_classes=n_classes)
  
  
  def make_layers(cfg, in_channels=3, batch_norm=False):
    layers = []
    for v in cfg:
      if v == 'M':
        layers += [nn.MaxPool2d(kernel_size=2, stride=2)]
      else:
        conv2d = nn.Conv2d(in_channels, v, kernel_size=3, padding=1)
        if batch_norm:
          layers += [conv2d, nn.BatchNorm2d(v), nn.ReLU(inplace=True)]
        else:
          layers += [conv2d, nn.ReLU(inplace=True)]
        in_channels = v
    return nn.Sequential(*layers)
  
  
  def get_vgg(in_channels=3, **kwargs):
    model = VGG(make_layers(cfg['D'], in_channels), **kwargs)
    return model
  