text: import functools
  
  import torch
  import torch.nn as nn
  from torchvision.models.densenet import _DenseLayer
  
  """
  Original code for ResnetGenerator and ResnetBlock provided by Yhu et al,
  taken from: https://github.com/junyanz/pytorch-CycleGAN-and-pix2pix
  
  See also:
  
  @inproceedings{CycleGAN2017,
    title={Unpaired Image-to-Image Translation using Cycle-Consistent Adversarial Networkss},
    author={Zhu, Jun-Yan and Park, Taesung and Isola, Phillip and Efros, Alexei A},
    booktitle={Computer Vision (ICCV), 2017 IEEE International Conference on},
    year={2017}
  }
  """
  
  
  class ResnetGenerator(nn.Module):
      """Resnet-based generator that consists of Resnet blocks between a few downsampling/upsampling operations.
      We adapt Torch code and idea from Justin Johnson's neural style transfer project
      (https://github.com/jcjohnson/fast-neural-style)
      """
  
      def __init__(self, input_nc, output_nc, ngf=64, norm_layer=nn.BatchNorm2d, use_dropout=False, n_blocks=6,
                   padding_type='reflect'):
          """Construct a Resnet-based generator
          Parameters:
              input_nc (int)      -- the number of channels in input images
              output_nc (int)     -- the number of channels in output images
              ngf (int)           -- the number of filters in the last conv layer
              norm_layer          -- normalization layer
              use_dropout (bool)  -- if use dropout layers
              n_blocks (int)      -- the number of ResNet blocks
              padding_type (str)  -- the name of padding layer in conv layers: reflect | replicate | zero
          """
          assert (n_blocks >= 0)
          super(ResnetGenerator, self).__init__()
          if type(norm_layer) == functools.partial:
              use_bias = norm_layer.func == nn.InstanceNorm2d
          else:
              use_bias = norm_layer == nn.InstanceNorm2d
  
          model = [nn.ReflectionPad2d(3),
                   nn.Conv2d(input_nc, ngf, kernel_size=7, padding=0, bias=use_bias),
                   norm_layer(ngf),
                   nn.ReLU(True)]
  
          n_downsampling = 2
          for i in range(n_downsampling):  # add downsampling layers
              mult = 2 ** i
              model += [nn.Conv2d(ngf * mult, ngf * mult * 2, kernel_size=3, stride=2, padding=1, bias=use_bias),
                        norm_layer(ngf * mult * 2),
                        nn.ReLU(True)]
  
          mult = 2 ** n_downsampling
          for i in range(n_blocks):  # add ResNet blocks
  
              model += [ResnetBlock(ngf * mult, padding_type=padding_type, norm_layer=norm_layer, use_dropout=use_dropout,
                                    use_bias=use_bias)]
  
          for i in range(n_downsampling):  # add upsampling layers
              mult = 2 ** (n_downsampling - i)
              model += [nn.ConvTranspose2d(ngf * mult, int(ngf * mult / 2),
                                           kernel_size=3, stride=2,
                                           padding=1, output_padding=1,
                                           bias=use_bias),
                        norm_layer(int(ngf * mult / 2)),
                        nn.ReLU(True)]
          model += [nn.ReflectionPad2d(3)]
          model += [nn.Conv2d(ngf, output_nc, kernel_size=7, padding=0)]
          # model += [nn.Tanh()]
          model += [nn.Sigmoid()]
  
          self.model = nn.Sequential(*model)
  
      def forward(self, input):
          """Standard forward"""
          return self.model(input)
  
  
  class ResnetBlock(nn.Module):
      """Define a Resnet block"""
  
      def __init__(self, dim, padding_type, norm_layer, use_dropout, use_bias):
          """Initialize the Resnet block
          A resnet block is a conv block with skip connections
          We construct a conv block with build_conv_block function,
          and implement skip connections in <forward> function.
          Original Resnet paper: https://arxiv.org/pdf/1512.03385.pdf
          """
          super(ResnetBlock, self).__init__()
          self.conv_block = self.build_conv_block(dim, padding_type, norm_layer, use_dropout, use_bias)
  
      def build_conv_block(self, dim, padding_type, norm_layer, use_dropout, use_bias):
          """Construct a convolutional block.
          Parameters:
              dim (int)           -- the number of channels in the conv layer.
              padding_type (str)  -- the name of padding layer: reflect | replicate | zero
              norm_layer          -- normalization layer
              use_dropout (bool)  -- if use dropout layers.
              use_bias (bool)     -- if the conv layer uses bias or not
          Returns a conv block (with a conv layer, a normalization layer, and a non-linearity layer (ReLU))
          """
          conv_block = []
          p = 0
          if padding_type == 'reflect':
              conv_block += [nn.ReflectionPad2d(1)]
          elif padding_type == 'replicate':
              conv_block += [nn.ReplicationPad2d(1)]
          elif padding_type == 'zero':
              p = 1
          else:
              raise NotImplementedError('padding [%s] is not implemented' % padding_type)
  
          conv_block += [nn.Conv2d(dim, dim, kernel_size=3, padding=p, bias=use_bias), norm_layer(dim), nn.ReLU(True)]
          if use_dropout:
              conv_block += [nn.Dropout(0.5)]
  
          p = 0
          if padding_type == 'reflect':
              conv_block += [nn.ReflectionPad2d(1)]
          elif padding_type == 'replicate':
              conv_block += [nn.ReplicationPad2d(1)]
          elif padding_type == 'zero':
              p = 1
          else:
              raise NotImplementedError('padding [%s] is not implemented' % padding_type)
          conv_block += [nn.Conv2d(dim, dim, kernel_size=3, padding=p, bias=use_bias), norm_layer(dim)]
  
          return nn.Sequential(*conv_block)
  
      def forward(self, x):
          """Forward function (with skip connections)"""
          out = x + self.conv_block(x)  # add skip connections
          return out
  
  
  class DenseBlock(nn.ModuleDict):
      """
      Based on torchvision.models.densenet, extended with an option to specify the normalisations for norm1 and norm2
      """
      _version = 2
  
      def __init__(self, num_layers, num_input_features, bn_size, growth_rate, drop_rate, memory_efficient=False,
                   norm_layer=nn.BatchNorm2d):
          super(DenseBlock, self).__init__()
          for i in range(num_layers):
              layer = _DenseLayer(
                  num_input_features + i * growth_rate,
                  growth_rate=growth_rate,
                  bn_size=bn_size,
                  drop_rate=drop_rate,
                  memory_efficient=memory_efficient,
              )
              if norm_layer != nn.BatchNorm2d:
                  layer.norm1 = norm_layer(num_input_features)
                  layer.norm2 = norm_layer(bn_size * growth_rate)
              self.add_module('denselayer%d' % (i + 1), layer)
  
      def forward(self, init_features):
          features = [init_features]
          for name, layer in self.items():
              new_features = layer(features)
              features.append(new_features)
          return torch.cat(features, 1)
  
  
  class DenseGenerator(nn.Module):
      """
      Based on ResnetGenerator, provided by Yhu et al, taken from: https://github.com/junyanz/pytorch-CycleGAN-and-pix2pix
      Modification replaces resnet blocks with dense ones.
      """
  
      def __init__(self, input_nc, output_nc, ngf=64, norm_layer=nn.BatchNorm2d, use_dropout=False, n_blocks=6,
                   padding_type='reflect'):
          """Construct a Resnet-based generator
          Parameters:
              input_nc (int)      -- the number of channels in input images
              output_nc (int)     -- the number of channels in output images
              ngf (int)           -- the number of filters in the last conv layer
              norm_layer          -- normalization layer
              use_dropout (bool)  -- if use dropout layers
              n_blocks (int)      -- the number of ResNet blocks
              padding_type (str)  -- the name of padding layer in conv layers: reflect | replicate | zero
          """
          assert (n_blocks >= 0)
          super(DenseGenerator, self).__init__()
          if type(norm_layer) == functools.partial:
              use_bias = norm_layer.func == nn.InstanceNorm2d
          else:
              use_bias = norm_layer == nn.InstanceNorm2d
  
          model = [nn.ReflectionPad2d(3),
                   nn.Conv2d(input_nc, ngf, kernel_size=7, padding=0, bias=use_bias),
                   norm_layer(ngf),
                   nn.ReLU(True)]
  
          n_downsampling = 2
          for i in range(n_downsampling):  # add downsampling layers
              mult = 2 ** i
              model += [nn.Conv2d(ngf * mult, ngf * mult * 2, kernel_size=3, stride=2, padding=1, bias=use_bias),
                        norm_layer(ngf * mult * 2),
                        nn.ReLU(True)]
  
          mult = 2 ** n_downsampling
          dense_features = ngf * mult
          dense_features = dense_features + 6 * 32
          for i in range(n_blocks):
              model += [DenseBlock(num_layers=6, num_input_features=ngf * mult, bn_size=4, growth_rate=32, drop_rate=0,
                                   norm_layer=norm_layer)]
              model += [norm_layer(dense_features), nn.ReLU(inplace=True),
                        nn.Conv2d(dense_features, ngf * mult, kernel_size=1, stride=1, bias=use_bias),
                        ]
  
          for i in range(n_downsampling):  # add upsampling layers
              mult = 2 ** (n_downsampling - i)
              model += [nn.ConvTranspose2d(ngf * mult, int(ngf * mult / 2),
                                           kernel_size=3, stride=2,
                                           padding=1, output_padding=1,
                                           bias=use_bias),
                        norm_layer(int(ngf * mult / 2)),
                        nn.ReLU(True)]
          model += [nn.ReflectionPad2d(3)]
          model += [nn.Conv2d(ngf, output_nc, kernel_size=7, padding=0)]
          model += [nn.Sigmoid()]
  
          self.model = nn.Sequential(*model)
  
      def forward(self, input):
          """Standard forward"""
          return self.model(input)
  