text: import torch
  import torch.nn as nn
  import torch.nn.functional as F
  
  import json
  from easydict import EasyDict as edict
  from graphs.models.custom_functions.weight_norm import WN_Conv3d, WN_ConvTranspose3d
  
  # 3D-UNet as Discriminator
  class Discriminator(nn.Module):
      def __init__(self, config):
          super(Discriminator, self).__init__()
          self.config = config
          self.input_channel = self.config.num_modalities
          self.num_classes = self.config.num_classes
          kernel_size = (3,3,3)
          kernel_size_deconv = (2,2,2)
          stride_deconv = (2,2,2)
          out_channels = 32
          self.lrelu = nn.LeakyReLU(0.2)
          self.dropout = nn.Dropout3d(p=0.2, inplace=False)
          self.final_activation = nn.Softmax(dim=1)
  
          # Defining the convolutional operations
          self.pool = nn.AvgPool3d(2)
  
          self.encoder0 = WN_Conv3d(self.input_channel, out_channels, kernel_size)
          self.encoder1 = WN_Conv3d(out_channels, out_channels, kernel_size)
  
          self.encoder2 = WN_Conv3d(out_channels, out_channels*(2), kernel_size)
          self.encoder3 = WN_Conv3d(out_channels*(2), out_channels*(2), kernel_size)
  
          self.encoder4 = WN_Conv3d(out_channels*(2), out_channels*(2**2), kernel_size)
          self.encoder5 = WN_Conv3d(out_channels*(2**2), out_channels*(2**2), kernel_size)
  
          self.encoder6 = WN_Conv3d(out_channels*(2**2), out_channels*(2**3), kernel_size)
          self.encoder7 = WN_Conv3d(out_channels*(2**3), out_channels*(2**3), kernel_size)
  
          self.decoder1 = WN_ConvTranspose3d(out_channels*(2**3), out_channels*(2**3), kernel_size_deconv, stride_deconv)
          #encoder5 + decoder1
          self.encoder8 = WN_Conv3d(out_channels*(2**2) + out_channels*(2**3), out_channels*(2**2), kernel_size)
          self.encoder9 = WN_Conv3d(out_channels*(2**2), out_channels*(2**2), kernel_size)
  
          self.decoder2 = WN_ConvTranspose3d(out_channels*(2**2), out_channels*(2**2), kernel_size_deconv, stride_deconv)
          #encoder3 + decoder2
          self.encoder10 = WN_Conv3d(out_channels*(2) + out_channels*(2**2), out_channels*(2), kernel_size)
          self.encoder11 = WN_Conv3d(out_channels*(2), out_channels*(2), kernel_size)
  
          self.decoder3 = WN_ConvTranspose3d(out_channels*(2), out_channels*(2), kernel_size_deconv, stride_deconv)
          #encoder1 + decoder3
          self.encoder12 = WN_Conv3d(out_channels + out_channels*(2), out_channels, kernel_size)
          self.encoder13 = WN_Conv3d(out_channels, out_channels, kernel_size)
  
          self.final_conv = WN_Conv3d(out_channels, self.num_classes, kernel_size)
  
      def forward(self, input, get_feature=False, use_dropout=False):
          conv0 = self.lrelu(self.encoder0(input))
          conv1 = self.lrelu(self.encoder1(conv0))
          pool1 = self.pool(conv1)
  
          conv2 = self.lrelu(self.encoder2(pool1))
          conv3 = self.lrelu(self.encoder3(conv2))
          pool3 = self.pool(conv3)
  
          conv4 = self.lrelu(self.encoder4(pool3))
          conv5 = self.lrelu(self.encoder5(conv4))
          pool5 = self.pool(conv5)
  
          conv6 = self.lrelu(self.encoder6(pool5))
          conv7 = self.lrelu(self.encoder7(conv6))
  
          if use_dropout:
              conv7 = self.dropout(conv7)
  
          deconv1 = self.decoder1(conv7)
          skip_connection1 = torch.cat((conv5, deconv1), 1)
          conv8 = self.lrelu(self.encoder8(skip_connection1))
          conv9 = self.lrelu(self.encoder9(conv8))
  
          deconv2 = self.decoder2(conv9)
          skip_connection2 = torch.cat((conv3, deconv2), 1)
          conv10 = self.lrelu(self.encoder10(skip_connection2))
          conv11 = self.lrelu(self.encoder11(conv10))
  
          deconv3 = self.decoder3(conv11)
          skip_connection3 = torch.cat((conv1, deconv3), 1)
          conv12 = self.lrelu(self.encoder12(skip_connection3))
          conv13 = self.lrelu(self.encoder13(conv12))
  
          if use_dropout:
              conv13 = self.dropout(conv13)
  
          final_output = self.final_conv(conv13)
  
          if not get_feature:
              return final_output, self.final_activation(final_output)
          else:
              return final_output, self.final_activation(final_output), conv6
  
  
  """
  netD testing
  """
  def main():
      config = json.load(open('../../configs/gan_exp_0.json'))
      config = edict(config)
      inp  = torch.autograd.Variable(torch.randn(config.batch_size, config.input_channels, config.patch_shape[0], config.patch_shape[1], config.patch_shape[2]))
      print (inp.shape)
      netD = Discriminator(config)
      out = netD(inp)
      print (out)
  
  if __name__ == '__main__':
      main()
  