text: # Copyright 2021 Dakewe Biotech Corporation. All Rights Reserved.
  # Licensed under the Apache License, Version 2.0 (the "License");
  #   you may not use this file except in compliance with the License.
  #   You may obtain a copy of the License at
  #
  #       http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  # ==============================================================================
  import torch
  import torch.nn as nn
  
  
  class DiscriminatorForVGG(nn.Module):
      def __init__(self, image_size: int = 128) -> None:
          super(DiscriminatorForVGG, self).__init__()
  
          feature_map_size = int(image_size // 16)
  
          self.features = nn.Sequential(
              nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1),  # input is (3) x 96 x 96
              nn.LeakyReLU(negative_slope=0.2, inplace=True),
  
              nn.Conv2d(64, 64, kernel_size=3, stride=2, padding=1, bias=False),  # state size. (64) x 48 x 48
              nn.BatchNorm2d(64),
              nn.LeakyReLU(negative_slope=0.2, inplace=True),
  
              nn.Conv2d(64, 128, kernel_size=3, stride=1, padding=1, bias=False),
              nn.BatchNorm2d(128),
              nn.LeakyReLU(negative_slope=0.2, inplace=True),
  
              nn.Conv2d(128, 128, kernel_size=3, stride=2, padding=1, bias=False),  # state size. (128) x 24 x 24
              nn.BatchNorm2d(128),
              nn.LeakyReLU(negative_slope=0.2, inplace=True),
  
              nn.Conv2d(128, 256, kernel_size=3, stride=1, padding=1, bias=False),
              nn.BatchNorm2d(256),
              nn.LeakyReLU(negative_slope=0.2, inplace=True),
  
              nn.Conv2d(256, 256, kernel_size=3, stride=2, padding=1, bias=False),  # state size. (256) x 12 x 12
              nn.BatchNorm2d(256),
              nn.LeakyReLU(negative_slope=0.2, inplace=True),
  
              nn.Conv2d(256, 512, kernel_size=3, stride=1, padding=1, bias=False),
              nn.BatchNorm2d(512),
              nn.LeakyReLU(negative_slope=0.2, inplace=True),
  
              nn.Conv2d(512, 512, kernel_size=3, stride=2, padding=1, bias=False),  # state size. (512) x 6 x 6
              nn.BatchNorm2d(512),
              nn.LeakyReLU(negative_slope=0.2, inplace=True)
          )
  
          self.classifier = nn.Sequential(
              nn.Linear(512 * feature_map_size * feature_map_size, 1024),
              nn.LeakyReLU(negative_slope=0.2, inplace=True),
              nn.Linear(1024, 1),
              nn.Sigmoid()
          )
  
      def forward(self, x: torch.Tensor) -> torch.Tensor:
          out = self.features(x)
          out = torch.flatten(out, 1)
          out = self.classifier(out)
  
          return out
  
  
  def discriminator_for_vgg(image_size: int = 128) -> DiscriminatorForVGG:
      r"""GAN model architecture from the
      `"One weird trick..." <https://arxiv.org/abs/1609.04802>`_ paper.
      """
      model = DiscriminatorForVGG(image_size)
      return model
  