<jupyter_start><jupyter_text><jupyter_code>import os
import random
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.nn.parallel
import torch.backends.cudnn as cudnn
import torch.optim as optim
import torch.utils.data
import torchvision.datasets as dset
import torchvision.transforms as transforms
import torchvision.utils as vutils
import torchvision.models as models
import numpy as np
import matplotlib.pyplot as plt
from tqdm import tqdm
import copy
from PIL import Image
from torch.autograd import Variable
import time
import warnings
from matplotlib.ticker import FuncFormatter
import itertools
 
warnings.filterwarnings('ignore')
 
%matplotlib inline
manual_seed = int(time.time())
print("Random Seed: ", manual_seed)
random.seed(manual_seed)
torch.manual_seed(manual_seed)
 
# location of folder
dataroot = "drive/MyDrive/SuperFolderTrainData"
# number of workers
workers = 2
# batch size for training
batch_size = 30
# image size for input
image_size = 128
# number of channel (1 for BW, 3 for RGB)
nc = 3
# Size of z latent vector (i.e. size of generator input)
nz = 16
# Size of feature maps in generator
ngf = 128
# Size of feature maps in discriminator
ndf = 128
# Number of training epochs
num_epochs = 50
# Learning rate for optimizers
lr = 0.0005
# Style loss weight
Msl_loss_weight = 0.003
SRsl_loss_weight = 0.015
# Colapse loss weight
Mcl_loss_weight = 0.015
SRcl_loss_weight = 0.002
# Beta values hyperparam for Adam optimizers
beta1 = 0.5
beta2 = 0.99
# Number of GPUs available, currently running on CPU
ngpu = 1
 
print(torch.cuda.get_device_name(torch.cuda.current_device()))
# define weights for layer and normalisation
def weights_init(m):
    classname = m.__class__.__name__
    if classname.find('Conv')!=-1:
        nn.init.normal_(m.weight.data, 0.0, 0.02)
    elif classname.find('BatchNorm')!=-1:
        nn.init.normal_(m.weight, 1.0, 0.02)
        nn.init.constant_(m.bias.data, 0)
 
# setup dataset and data loader
# dataloader for microstructure images
Mdataset = dset.ImageFolder(root=dataroot, transform=transforms.Compose([transforms.Resize(int(image_size/2)), 
                                                                        transforms.CenterCrop(int(image_size/2)),
                                                                        transforms.ToTensor(),
                                                                        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])]))
 
Mdataloader = torch.utils.data.DataLoader(Mdataset, batch_size=batch_size, shuffle=True, num_workers=workers)

# dataloader for sr images
SRdataset = dset.ImageFolder(root=dataroot, transform=transforms.Compose([transforms.Resize(image_size), 
                                                                        transforms.CenterCrop(image_size),
                                                                        transforms.ToTensor(),
                                                                        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])]))

SRdataloader = torch.utils.data.DataLoader(SRdataset, batch_size=batch_size, shuffle=True, num_workers=workers)

 
device = torch.device("cuda:0" if(torch.cuda.is_available() and ngpu>0) else "cpu")<jupyter_output><empty_output><jupyter_text>**Generator Module**<jupyter_code>class Generator(nn.Module):
  def __init__(self, ngpu):
    super(Generator, self).__init__()
    self.ngpu = ngpu
    self.main = nn.Sequential()
    
    # input z into convolution
    self.main.add_module('conv1', nn.ConvTranspose2d(nz, ngf*8, 8, 2, 0, bias=False))
    self.main.add_module('btch1', nn.BatchNorm2d(ngf*8))
    self.main.add_module('relu1', nn.ReLU(True))

    # state ngf * 8 * 8
    self.main.add_module('conv2', nn.ConvTranspose2d(ngf*8, ngf*4, 4, 2, 1, bias=False))
    self.main.add_module('btch2', nn.BatchNorm2d(ngf*4))
    self.main.add_module('relu2', nn.ReLU(True))
            
    # state ngf/2 * 16*16
    self.main.add_module('conv3', nn.ConvTranspose2d(ngf*4, ngf*2, 4, 2, 1, bias=False))
    self.main.add_module('btch3', nn.BatchNorm2d(ngf*2))
    self.main.add_module('relu3', nn.ReLU(True))
            
    # state ngf/4 * 32*32
    self.main.add_module('conv4', nn.ConvTranspose2d(ngf*2, ngf, 4, 2, 1, bias=False))
    self.main.add_module('btch4', nn.BatchNorm2d(ngf))
    self.main.add_module('relu4', nn.ReLU(True))
            
    # state ngf/8 *64*64
    self.main.add_module('conv5', nn.ConvTranspose2d(ngf, nc, 1, 1, 0, bias=False))
    self.main.add_module('output', nn.Tanh())
        
    #state 3 *64*64
        
  def forward(self, input):
    return self.main(input)
# Create Generator for producing microstructure images
mGen = Generator(ngpu=ngpu).to(device)
mGen.apply(weights_init)

print(mGen)
class DenseLayer(nn.Module):
  def __init__(self, in_size, growth_rate, bn_size, drop_rate=0.0):
    super(DenseLayer, self).__init__()
    self.bottleneck = nn.Sequential() # define bottleneck layers
    self.bottleneck.add_module('btch1', nn.BatchNorm2d(in_size))
    self.bottleneck.add_module('relu1', nn.ReLU(inplace=True))
    self.bottleneck.add_module('conv1', nn.ConvTranspose2d(in_size, bn_size*growth_rate, kernel_size=1, stride=1, padding=0, bias=False))
 
    self.basic = nn.Sequential() # define basic block
    self.basic.add_module('btch2', nn.BatchNorm2d(bn_size*growth_rate))
    self.basic.add_module('relu2', nn.ReLU(inplace=True))
    self.basic.add_module('conv2', nn.ConvTranspose2d(bn_size*growth_rate, growth_rate, kernel_size=3, stride=1, padding=1, bias=False))
 
    self.droprate = drop_rate
 
  def forward(self, input):
    out = self.bottleneck(input)
    if self.droprate > 0:
      out = F.dropout(out, p=self.droprate, inplace=False, training=self.training)
    
    out = self.basic(out)
    if self.droprate > 0:
      out = F.dropout(out, p=self.droprate, inplace=False, training=self.training)
    return torch.cat((input,out), 1)
 
class TransitionLayer(nn.Module):
  def __init__(self, in_size, out_size, drop_rate=0.0):
    super(TransitionLayer, self).__init__()
    self.trans = nn.Sequential()
    self.trans.add_module('btch', nn.BatchNorm2d(in_size))
    self.trans.add_module('relu', nn.ReLU(inplace=True))
    self.trans.add_module('conv', nn.ConvTranspose2d(in_size, out_size, kernel_size=1, stride=1, padding=0, bias=False))
 
    self.droprate = drop_rate
 
  def forward(self, input):
    out = self.trans(input)
    if self.droprate > 0:
      out = F.dropout(out, p=self.droprate, inplace=False, training=self.training)
 
    return out

class DenseBlock(nn.Module):
  def __init__(self, num_layers, in_size, growth_rate, bn_size, droprate=0.0):
    super(DenseBlock, self).__init__()
    self.dblayer = nn.Sequential()
 
    for i in range(num_layers):
      layer = DenseLayer(in_size=(in_size+i*growth_rate), growth_rate=growth_rate, bn_size=bn_size, drop_rate=droprate)
      self.dblayer.add_module('denselayer{}'.format(i+1), layer)
 
  def forward(self, input):
    return self.dblayer(input)

class GenDenseNet(nn.Module):
  def __init__(self, ngpu, growth_rate=12, block_config=(1,2,2,1), in_size=512, bn_size=4, drop_rate=0.0):
    super(GenDenseNet, self).__init__()
    self.ngpu = ngpu
    self.features = nn.Sequential()
    self.features.add_module('btch0', nn.BatchNorm2d(in_size))
 
    num_features = in_size
    for i, num_layers in enumerate(block_config):
      block = DenseBlock(num_layers=num_layers, in_size=num_features, growth_rate=growth_rate, bn_size=bn_size, droprate=drop_rate)
      self.features.add_module('denseblock{}'.format(i+1), block)
      num_features += num_layers*growth_rate
 
      if i!=len(block_config)-1:
        trans = TransitionLayer(in_size=num_features, out_size=num_features*2, drop_rate=drop_rate)
        self.features.add_module('transitionblock{}'.format(i+1), trans)
        num_features *= 2
 
    self.features.add_module('convfinal', nn.ConvTranspose2d(num_features, 3, kernel_size=8, stride=2, padding=3, bias=False))
    self.features.add_module('Tanh', nn.Tanh())
 
  def forward(self, input):
    return self.features(input)
# Create Generator for producing SR images
srGen = GenDenseNet(ngpu=ngpu, in_size=3).to(device)
srGen.apply(weights_init)

print(srGen)<jupyter_output>GenDenseNet(
  (features): Sequential(
    (btch0): BatchNorm2d(3, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
    (denseblock1): DenseBlock(
      (dblayer): Sequential(
        (denselayer1): DenseLayer(
          (bottleneck): Sequential(
            (btch1): BatchNorm2d(3, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
            (relu1): ReLU()
            (conv1): ConvTranspose2d(3, 48, kernel_size=(1, 1), stride=(1, 1), bias=False)
          )
          (basic): Sequential(
            (btch2): BatchNorm2d(48, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
            (relu2): ReLU()
            (conv2): ConvTranspose2d(48, 12, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
          )
        )
      )
    )
    (transitionblock1): TransitionLayer(
      (trans): Sequential(
        (btch): BatchNorm2d(15, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        (relu): ReLU()
        (conv): [...]<jupyter_text>**Discriminator Module**<jupyter_code>class Discriminator(nn.Module):
  def __init__(self, ngpu, final_kernel_size=8): #final_kernel_size=8 ->128 | =4 -> 64 image size
    super(Discriminator, self).__init__()
    self.ngpu = ngpu
    self.main = nn.Sequential()

    # input is (nc) *128*128
    self.main.add_module('conv1', nn.Conv2d(nc, ndf, kernel_size=4, stride=2, padding=1, bias=False))
    self.main.add_module('btch1', nn.BatchNorm2d(ndf))
    self.main.add_module('relu1', nn.LeakyReLU(0.2, inplace=True))

    # input is ndf/8 *64*64
    self.main.add_module('conv2', nn.Conv2d(ndf, ndf*2, kernel_size=4, stride=2, padding=1, bias=False))
    self.main.add_module('btch2', nn.BatchNorm2d(ndf*2))
    self.main.add_module('relu2', nn.LeakyReLU(0.2, inplace=True))

    #input is ndf/4 *32*32
    self.main.add_module('conv3', nn.Conv2d(ndf*2, ndf*4, kernel_size=4, stride=2, padding=1, bias=False))
    self.main.add_module('btch3', nn.BatchNorm2d(ndf*4))
    self.main.add_module('relu3', nn.LeakyReLU(0.2, inplace=True))

    #input is ndf/2 *16*16
    self.main.add_module('conv4', nn.Conv2d(ndf*4, ndf*8, kernel_size=4, stride=2, padding=1, bias=False))
    self.main.add_module('btch4', nn.BatchNorm2d(ndf*8))
    self.main.add_module('relu4', nn.LeakyReLU(0.2, inplace=True))

    #input is ndf *8*8
    self.main.add_module('conv5', nn.Conv2d(ndf*8, 1, kernel_size=final_kernel_size, stride=1, padding=0, bias=False))
    self.main.add_module('output', nn.Sigmoid())

  def forward(self, input):
    return self.main(input)
# Create discriminitor to distinguish microstructure
mDis = Discriminator(ngpu=ngpu, final_kernel_size=4).to(device)
mDis.apply(weights_init)

print(mDis)
# Create discriminator to distinguish SR images
srDis = Discriminator(ngpu=ngpu).to(device)
srDis.apply(weights_init)

print(srDis)<jupyter_output>Discriminator(
  (main): Sequential(
    (conv1): Conv2d(3, 128, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1), bias=False)
    (btch1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
    (relu1): LeakyReLU(negative_slope=0.2, inplace=True)
    (conv2): Conv2d(128, 256, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1), bias=False)
    (btch2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
    (relu2): LeakyReLU(negative_slope=0.2, inplace=True)
    (conv3): Conv2d(256, 512, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1), bias=False)
    (btch3): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
    (relu3): LeakyReLU(negative_slope=0.2, inplace=True)
    (conv4): Conv2d(512, 1024, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1), bias=False)
    (btch4): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
    (relu4): LeakyReLU(negative_slope=0.2, i[...]<jupyter_text>Check for inputs and outputs consistency<jupyter_code>in_tensor = torch.rand(10, 16, 1, 1).to(device)

gen_img = mGen(in_tensor)
print(gen_img.shape)
sr_img = srGen(gen_img)
print(sr_img.shape)
mdis_tensor = mDis(gen_img)
print(mdis_tensor.shape)
srdis_tensor = srDis(sr_img)
print(srdis_tensor.shape)<jupyter_output>torch.Size([10, 1, 1, 1])<jupyter_text>**Error Functions and Optimizers**  Style Loss<jupyter_code>def gram_matrix(input):
  a, b, c, d = input.size() # a=batch size(=1) b=number of feature maps (c,d)=dimensions of a f. map (N=c*d)
  features = input.view(a * b, c * d) # resise F_XL into \hat F_XL
  G = torch.mm(features, features.t()) # compute the gram product
  return G.div(a * b * c * d) # we 'normalize' the values of the gram matrix by dividing by the number of element in each feature maps.

def style_loss(style_layer, combination_layer, num_channels = nc):
  assert style_layer.size()==combination_layer.size(), "Input Sizes do not match in StyleLoss"
  a, b, c, d = style_layer.size() # a=batch size(=1) b=number of feature maps (c,d)=dimensions of a f. map (N=c*d)
  styleloss = 0
  
  for i in range(a):
    S = gram_matrix(torch.unsqueeze(style_layer[i], 0)).detach()
    C = gram_matrix(torch.unsqueeze(combination_layer[i],0)).detach()
    styleloss = torch.add(styleloss, torch.sum(torch.square(S-C)/(2*num_channels*c*d)**2)*3e6)
    styleloss = Variable(styleloss, requires_grad=True)

  return styleloss

def content_loss(content_layer, combination_layer, num_channels = nc):
  assert content_layer.size()==combination_layer.size(), "Input Sizes do not match in Content loss"
  a, b, c, d = content_layer.size()
  ctloss = torch.sum(torch.square(content_layer-combination_layer)/(2*num_channels*c*d)**2)
  ctloss = Variable(ctloss, requires_grad=True)
  return ctloss

def colapse_loss(target_feature):
  z_d_gen = torch.flatten(target_feature).reshape(1,-1)
  nom = torch.mm(z_d_gen, z_d_gen.t())
  denom = torch.sqrt(torch.sum(torch.square(z_d_gen), 1, keepdim=True))
  pt = torch.square(torch.transpose((nom / denom), 1, 0))
  #pt = pt - torch.diag(torch.diag(pt))
  pulling_term = torch.sum(pt) / (batch_size * (batch_size - 1)*4e1)
  pulling_term = Variable(pulling_term, requires_grad=True)

  return pulling_term

vgg19 = models.vgg19(pretrained=True).features.to(device).eval()

def get_vgg_layers(input, model=vgg19, num_layer=4):
  i = 0
  Model = copy.deepcopy(model)
  list_layers = nn.ModuleList()
  module = None
  for layer in Model.children():
    if isinstance(layer, nn.Conv2d):
      i +=1
      name = 'conv{}'.format(i)
      list_layers.append(module)
      module = nn.Sequential()
    elif isinstance(layer, nn.ReLU):
      name = 'relu{}'.format(i)
    elif isinstance(layer, nn.MaxPool2d):
      name = 'pool{}'.format(i)
    elif isinstance(layer, nn.BatchNorm2d):
      name = 'btnm{}'.format(i)
    else: 
      raise RuntimeError('Unrecognized layer: {}'.format(layer.__class__.__name__))

    module.add_module(name, layer)

    if i >= (num_layer+1):
      break

  outconv1 = list_layers[1](input)
  outconv2 = list_layers[2](outconv1)
  outconv3 = list_layers[3](outconv2)
  outconv4 = list_layers[4](outconv3)

  return outconv1, outconv2, outconv3, outconv4
def get_style_collapse_loss(style_batch, gen_batch):
  Sconv1, Sconv2, Sconv3, Sconv4 = get_vgg_layers(style_batch)
  Cconv1, Cconv2, Cconv3, Cconv4 = get_vgg_layers(gen_batch)
  
  # style loss
  s1 = style_loss(Sconv1, Cconv1)
  s2 = style_loss(Sconv2, Cconv2)
  s3 = style_loss(Sconv3, Cconv3)
  s4 = style_loss(Sconv4, Cconv4)
  sl_loss = torch.mean(s1+s2+s3+s4)
  
  Cconv = torch.cat((torch.flatten(Cconv1), torch.flatten(Cconv2), torch.flatten(Cconv3), torch.flatten(Cconv4)), 0)
  c = colapse_loss(Cconv)
  cl_loss = torch.mean(c)
  ct_loss = content_loss(Sconv4, Cconv4)
  return sl_loss, cl_loss, ct_loss<jupyter_output><empty_output><jupyter_text>Optimizers and BCE loss<jupyter_code># Initialize BCELoss function
criterion = nn.BCELoss()
 
# Create batch of latent vectors that we will use to visualize
#  the progression of the generator
fixed_noise = torch.randn(1, 16, 1, 1, device=device)
 
# Establish convention for real and fake labels during training
real_label = 1.
fake_label = 0.
 
# Setup Adam optimizers for all networks
optimizerMD = optim.Adam(mDis.parameters(), lr=lr, betas=(0.9, 0.999), weight_decay=0.3)
optimizerSRD = optim.Adam(srDis.parameters(), lr=lr, betas=(0.9, 0.999), weight_decay=0.3)
optimizerMG = optim.Adam(mGen.parameters(), lr=lr, betas=(0.9, 0.999), weight_decay=0.3)
optimizerSRG = optim.SGD(srGen.parameters(), lr=lr, momentum=0.9, dampening=0, weight_decay=0.3, nesterov=True)<jupyter_output><empty_output><jupyter_text>**Training Loop**<jupyter_code>mimg_list = [] # stores image produced by mGen
srimg_list = [] # stores image produces by srGen
mGen_loss = [] # stores losses for mGen
srGen_loss = [] # stores losses for srGen
mDis_loss = [] # stores losses for mDis
srDis_loss = [] # stores losses for srDis
mGen_bce = []
srGen_bce = []
iter = 0 # keeps track of total iterations
skip_steps = 5

print('Starting Training')
for epoch in tqdm(range(num_epochs)):
  for (i, mdata), srdata in zip(enumerate(Mdataloader), SRdataloader):
    MDis_real, MDis_fake, SRDis_real, SRDis_fake = 0.0, 0.0, 0.0, 0.0
    errMDis_real, errMDis_fake, errMDis = None, None, None
    errSRDis_real, errSRDis_fake, errSRDis = None, None, None
    ########### Train Discriminators ############3
    if i%skip_steps == 0:
      ############ MDis track ###############
      ############# REAL DATA ###############
      mDis.zero_grad()
      real_mdis_input = mdata[0].to(device)
      b_size = real_mdis_input.size(0)
      label = torch.full((b_size,), real_label, dtype=torch.float, device=device)
      moutput = mDis(real_mdis_input).view(-1)

      errMDis_real = criterion(moutput, label)
      errMDis_real.backward()
      MDis_real = moutput.mean().item()

      ############# FAKE DATA ###############
      noise = torch.randn(b_size, 16, 1, 1, device=device)
      fake_mdis_input = mGen(noise)
      label.fill_(fake_label)
      moutput = mDis(fake_mdis_input.detach()).view(-1)

      errMDis_fake = criterion(moutput, label)
      errMDis_fake.backward()
      MDis_fake = moutput.mean().item()

      ############# OPTIMIZE MDIS ############
      errMDis = errMDis_real + errMDis_fake
      optimizerMD.step()

      ############ SRDis track ##############
      ############# REAL DATA ###############
      srDis.zero_grad()
      real_srdis_input = srdata[0].to(device)
      b_size = real_srdis_input.size(0)
      label.fill_(real_label)
      sroutput = srDis(real_srdis_input).view(-1)

      errSRDis_real = criterion(sroutput, label)
      errSRDis_real.backward()
      SRDis_real = sroutput.mean().item()

      ############# FAKE DATA ###############
      fake_srdis_input = srGen(fake_mdis_input)
      label.fill_(fake_label)
      sroutput = srDis(fake_srdis_input.detach()).view(-1)

      errSRDis_fake = criterion(sroutput, label)
      errSRDis_fake.backward()
      SRDis_fake = sroutput.mean().item()

      ############# OPTIMIZE SRDIS ############
      errSRDis = errSRDis_real + errSRDis_fake
      optimizerSRD.step()

    ############ Gen track ###############
    mGen.zero_grad()
    srGen.zero_grad()

    ############ MGen track ###############
    if i%skip_steps != 0:
      real_mdis_input = mdata[0].to(device)
      label = torch.full((b_size,), real_label, dtype=torch.float, device=device)
      noise = torch.randn(b_size, 16, 1, 1, device=device)
      fake_mdis_input = mGen(noise)
      moutput = mDis(fake_mdis_input.detach()).view(-1)

    else:
      label.fill_(real_label)
      moutput = mDis(fake_mdis_input.detach()).view(-1)

    errMGen_bce = criterion(moutput, label)
    errMGen_bce.backward()

    Msl_loss, Mcl_loss,_ = get_style_collapse_loss(real_mdis_input, fake_mdis_input)
    Msl_loss *= Msl_loss_weight
    Msl_loss.backward()

    Mcl_loss *= Mcl_loss_weight
    Mcl_loss.backward()

    errMGen = errMGen_bce + Msl_loss + Mcl_loss
    optimizerMG.step()

    ############ SRGen track ###############
    if i%skip_steps != 0:
      real_srdis_input = srdata[0].to(device)
      label.fill_(real_label)
      fake_srdis_input = srGen(fake_mdis_input)
      sroutput = srDis(fake_srdis_input.detach()).view(-1)

    else:
      label.fill_(real_label)
      sroutput = srDis(fake_srdis_input.detach()).view(-1)

    errSRGen_bce = criterion(sroutput, label)
    errSRGen_bce.backward()

    SRsl_loss, SRcl_loss, SRct_loss = get_style_collapse_loss(real_srdis_input, fake_srdis_input)

    SRct_loss.backward()
    SRsl_loss *= SRsl_loss_weight
    SRsl_loss.backward()

    SRcl_loss *= SRcl_loss_weight
    SRcl_loss.backward()
    errSRGen = errSRGen_bce + SRsl_loss + SRct_loss + SRcl_loss
    optimizerSRG.step()

    if i%skip_steps == 0:
      mGen_loss.append(errMGen.item())
      srGen_loss.append(errSRGen.item())
      mDis_loss.append(errMDis.item())
      srDis_loss.append(errSRDis.item())
      mGen_bce.append(errMGen_bce.item())
      srGen_bce.append(errSRGen_bce.item())

    if iter%500 == 0:
      with torch.no_grad():
        fakem = mGen(fixed_noise)
      mimg_list.append(fakem.detach().cpu())
      srimg_list.append(srGen(fakem).detach().cpu())
    
    iter += 1

    if i%60 == 0:
      print("[{}/{}]".format(i, len(Mdataloader)))
      print("MDis_real: {:.4f} | MDis_fake: {:.4f} | SRDis_real: {:.4f} | SRDis_fake: {:.4f}".format(MDis_real, MDis_fake, SRDis_real, SRDis_fake))
      print("errMDis: {:.4f} | errMDis_real: {:.4f} | errMDis_fake: {:.4f}".format(errMDis.item(), errMDis_real.item(), errMDis_fake.item()))
      print("errSRDis: {:.4f} | errSRDis_real: {:.4f} | errSRDis_fake: {:.4f}".format(errSRDis.item(), errSRDis_real.item(), errSRDis_fake.item()))
      print("errMGen: {:.4f} | errMGen_bce: {:.4f} | Msl_loss: {:.4f} | Mcl_loss: {:.4f}".format(errMGen.item(), errMGen_bce.item(), Msl_loss.item(), Mcl_loss.item()))
      print("errSRGen: {:.4f} | errSRGen_bce: {:.4f}| SRsl_loss: {:.4f} | SRct_loss: {:.4f} | SRcl_loss: {:.4f}".format(errSRGen.item(), errSRGen_bce.item(), SRsl_loss.item(), SRct_loss.item(), SRcl_loss.item()))
  
  if epoch%25 == 0:
    print("Saving models at epoch: {}".format(epoch))
    torch.save(mGen.state_dict(), '/content/drive/MyDrive/mGen-{}.pt'.format(epoch))
    torch.save(mDis.state_dict(), '/content/drive/MyDrive/mDis-{}.pt'.format(epoch))
    torch.save(srGen.state_dict(), '/content/drive/MyDrive/srGen-{}.pt'.format(epoch))
    torch.save(srDis.state_dict(), '/content/drive/MyDrive/srDis-{}.pt'.format(epoch))

print("Training Finished")<jupyter_output>0%|          | 0/50 [00:00<?, ?it/s][A[A[A[A[A[A[A[A<jupyter_text>**Plotting Losses and Visualisation**<jupyter_code>plt.figure(figsize=(10,5))
plt.title("Generator and Discriminator Loss During Training")
plt.plot(mGen_loss,label="MGen")
plt.plot(srGen_loss,label="SRGen")
plt.plot(mDis_loss,label="MDis")
plt.plot(srDis_loss,label="SRDis")
plt.xlabel("iterations")
plt.ylabel("Loss")
plt.legend()
plt.gca().get_xaxis().set_major_formatter(FuncFormatter(lambda x, p: format(int(x*skip_steps), ',')))
plt.show()
plt.figure(figsize=(10,5))
plt.title("Generator BCE Loss During Training")
plt.plot(mGen_bce,label="MGenBCE")
plt.plot(srGen_bce,label="SRGenBCE")
plt.xlabel("iterations")
plt.ylabel("Loss")
plt.legend()
plt.gca().get_xaxis().set_major_formatter(FuncFormatter(lambda x, p: format(int(x*skip_steps), ',')))
plt.show()
plt.figure(figsize=(30,10))
for i in range(len(mimg_list)-10, len(mimg_list)):
  plt.subplot(2,10/2,i-(len(mimg_list)-10)+1)
  plt.imshow(mimg_list[i-(len(mimg_list)-10)].mean(0).mean(0))#, cmap='gray')
  plt.title('{}'.format((i-(len(mimg_list)-10))*500))
plt.show()
plt.figure(figsize=(30,10))
for i in range(len(srimg_list)-10,len(srimg_list)):
  plt.subplot(2,10/2,i-(len(srimg_list)-10)+1)
  plt.imshow(srimg_list[i-(len(srimg_list)-10)].mean(0).mean(0))#, cmap='gray')
  plt.title('{}'.format(i-(len(srimg_list)-10)*500))
plt.show()

print("Saving models at epoch: {}".format(num_epochs))
torch.save(mGen.state_dict(), '/content/drive/MyDrive/mGen-{}.pt'.format(num_epochs))
torch.save(mDis.state_dict(), '/content/drive/MyDrive/mDis-{}.pt'.format(num_epochs))
torch.save(srGen.state_dict(), '/content/drive/MyDrive/srGen-{}.pt'.format(num_epochs))
torch.save(srDis.state_dict(), '/content/drive/MyDrive/srDis-{}.pt'.format(num_epochs))<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.15684368291003362, 'fraction_numerical': 0.030015797788309637, 'mean_word_length': 4.175681341719078, 'pattern_counts': {'":': 0, '<': 30, '<?xml version=': 0, '>': 36, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 18, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '8', 'dir': 'jupyter-structured-clean-dedup', 'id': '546479', 'license': 'MIT', 'n_tokens_mistral': 9748, 'n_tokens_neox': 9200, 'n_words': 1640, 'path': 'EnsembleGAN.ipynb', 'repo_name': 'abhistar/DDP_GAN', 'size': '24687'}