BigGANModel -- Trying to train a model to generate images

[btb4198]

I am trying to train a model to generate images, I have a dataset of over 17K of men posing. I have been training my model for a few hours now and Sadly all I am getting is this:

Also my d_loss=-9.79e+3 how is it Negative ?

here is my code:

# -*- coding: utf-8 -*-
"""BigGANModel.ipynb

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/drive/1mw6J_dBCCmx6_mwpa7VDHNG7PiR3K2th
"""

import torch
import torch.nn as nn
import torchvision
import urllib.request
from torchvision.transforms import Resize
from torchvision.utils import save_image
from torchvision.transforms import ToPILImage
from torchvision import transforms, utils, datasets
from torch.utils.data import Dataset, DataLoader
import matplotlib.pyplot as plt
import numpy as np
from io import BytesIO
import torchvision.transforms as T
import torchvision.transforms.functional as F
from PIL import Image
from google.colab import drive
from google.colab import files
from tqdm import tqdm
import pickle
import random
import cv2
import os
import torch.optim as optim
from torch.autograd import Variable
drive.mount('/content/gdrive')

from sys import path
path.append("/content/gdrive/My Drive/Python_Libraries")
import genericgandataset as ggd

from sys import path
path.append("/content/gdrive/My Drive/Python_Libraries")
import reuseablecustompythonfunctions as rcpf

batch_size = 32
num_workers = 4
LEARNING_RATE = 2e-4
betas=(0.5, 0.999)
device = "cuda" if torch.cuda.is_available() else "cpu"
ggd.ImageHeight = 128
ggd.ImageWidth = 128
generator_ImageHeight = 1024
generator_ImageWidth = 1024
val_frequency = 100
evolution_images_per_epoch = 8
val_frequency = 2
LOAD_MODEL = True
SAVE_MODEL = True
CHECKPOINT_DISC = "gdrive/My Drive/All_Deep_Learning_Models/CNN_Models/BigGANModel/Disc_BigGANModel_"
CHECKPOINT_GEN = "gdrive/My Drive/All_Deep_Learning_Models/CNN_Models/BigGANModel/Gen_BigGANModel_"
LOAD_CHECKPOINT_DISC = "/content/gdrive/MyDrive/All_Deep_Learning_Models/CNN_Models/BigGANModel/Disc_BigGANModel_0.pth.tar"
LOAD_CHECKPOINT_GEN = "/content/gdrive/MyDrive/All_Deep_Learning_Models/CNN_Models/BigGANModel/Gen_BigGANModel_0.pth.tar"
evolution_folder ="gdrive/My Drive/All_Deep_Learning_Models/BigGANModelEvolution"

class ResBlock(nn.Module):
    def __init__(self, in_channels, out_channels, upsample=False, downsample=False):
        super(ResBlock, self).__init__()
        self.upsample = upsample
        self.downsample = downsample
        self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1)
        self.bn1 = nn.BatchNorm2d(out_channels)
        self.relu = nn.ReLU(inplace=False)
        self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1)
        self.bn2 = nn.BatchNorm2d(out_channels)

        self.shortcut = nn.Sequential()
        if upsample or downsample or in_channels != out_channels:
            self.shortcut = nn.Sequential(
                nn.Conv2d(in_channels, out_channels, kernel_size=1, padding=0),
                nn.BatchNorm2d(out_channels)
            )

    def forward(self, x):
        out = self.relu(self.bn1(self.conv1(x)))
        out = self.bn2(self.conv2(out))

        shortcut_x = self.shortcut(x)

        if self.upsample:
            out = nn.functional.interpolate(out, scale_factor=2)
            shortcut_x = nn.functional.interpolate(shortcut_x, scale_factor=2)
        if self.downsample:
            out = nn.functional.avg_pool2d(out, 2)
            shortcut_x = nn.functional.avg_pool2d(shortcut_x, 2)

        out += shortcut_x
        return self.relu(out)

class BigGANGenerator(nn.Module):
    def __init__(self, latent_dim=128):
        super(BigGANGenerator, self).__init__()
        self.latent_dim = latent_dim
        self.proj = nn.Linear(latent_dim, 16 * latent_dim * 4 * 4)
        self.bn = nn.BatchNorm1d(16 * latent_dim * 4 * 4)
        self.relu = nn.ReLU(inplace=True)

        self.res_blocks = nn.Sequential(
            ResBlock(16 * latent_dim, 8 * latent_dim, upsample=True),
            ResBlock(8 * latent_dim, 4 * latent_dim, upsample=True),
            ResBlock(4 * latent_dim, 2 * latent_dim, upsample=True),
            ResBlock(2 * latent_dim, latent_dim, upsample=True),
            ResBlock(latent_dim, latent_dim // 2, upsample=True),  # Add this layer
            ResBlock(latent_dim // 2, latent_dim // 4, upsample=True),  # Add this layer
        )

        self.conv = nn.Conv2d(latent_dim // 4, 3, kernel_size=3, padding=1)
        self.tanh = nn.Tanh()

    def forward(self, x):
        x = self.proj(x)
        x = self.bn(x)
        x = self.relu(x)
        x = x.view(-1, 16 * self.latent_dim, 4, 4)

        x = self.res_blocks(x)
        x = self.conv(x)
        x = self.tanh(x)
        return x

class BigGANDiscriminator(nn.Module):
    def __init__(self):
        super(BigGANDiscriminator, self).__init__()

        self.res_blocks = nn.Sequential(
            ResBlock(3, 64, downsample=True),
            ResBlock(64, 128, downsample=True),
            ResBlock(128, 256, downsample=True),
            ResBlock(256, 512, downsample=True),
            ResBlock(512, 1024, downsample=True),
        )

        self.relu = nn.ReLU(inplace=False)
        self.pool = nn.AdaptiveAvgPool2d(1)
        self.fc = nn.Linear(1024, 1)

    def forward(self, x):
        x = self.res_blocks(x)
        x = self.relu(x)
        x = self.pool(x)
        x = x.view(x.size(0), -1)
        x = self.fc(x)
        return x

# Initialize the generator and discriminator objects
generator = BigGANGenerator(latent_dim=128)
discriminator = BigGANDiscriminator()

# Move the generator and discriminator to GPU if available
generator = generator.to(device)
discriminator = discriminator.to(device)

# Set up the optimizers
optimizer_G = optim.Adam(generator.parameters(), lr=LEARNING_RATE , betas=betas)
optimizer_D = optim.Adam(discriminator.parameters(), lr=LEARNING_RATE , betas=betas)

image_saved_count = 0
g_scaler = torch.cuda.amp.GradScaler()
d_scaler = torch.cuda.amp.GradScaler()

train_loader = ggd.getdataloaders(batch_size, num_workers)

iterations_per_epoch = ggd.size_of_trainning_dataset // batch_size
evolution_frequency = iterations_per_epoch // evolution_images_per_epoch

if LOAD_MODEL:
  rcpf.load_checkpoint(LOAD_CHECKPOINT_GEN, generator, optimizer_G,LEARNING_RATE,)
  rcpf.load_checkpoint(LOAD_CHECKPOINT_DISC, discriminator, optimizer_D,LEARNING_RATE,)

def test_discriminator():
    ImageHeight, ImageWidth = 128, 128
    x = torch.randn((5, 3, ImageHeight, ImageWidth))
    test_model = BigGANDiscriminator()
    test_model = test_model.to(device)
    x = x.to(device)
    preds = test_model(x)
    print(preds.shape)

test_discriminator()

def test_generator():  
    # Create a Generator model
    test_generator = BigGANGenerator(latent_dim=128)
    test_generator = test_generator.to(device)
    # Create a random latent vector
    z = torch.randn(2, 128).to(device)  # or any batch size greater than 1
    # Generate an image using the generator
    img_generated = test_generator(z)
    # Print the generated image shape
    print(img_generated.shape)

test_generator()

def test_model(generator, device=device):
    generator.eval()  # Set the generator to evaluation mode
    with torch.no_grad():  # No need to compute gradients for this operation
        random_noise = torch.randn(1, 512).to(device)  # Create random noise vector
        generated_image = generator(random_noise)  # Generate image using the generator
        # Convert the resized image to a PIL image and display it
        image = ToPILImage()(generated_image)
        plt.imshow(image)
        plt.axis('off')
        plt.show()

# Define the loss function
def wasserstein_loss(output, target):
  return torch.mean(output * target)

def save_images(y_fake, folder, image_saved_count):
  save_image(y_fake, f"{folder}/Generated_Image{image_saved_count}.png")

import torch.nn.functional as F
def resize_images(images, size=128):
    # Check if the input images have the expected shape
    if len(images.shape) != 4:
        raise ValueError("Expected input shape: [batch_size, channels, height, width]")

    # Rescale the images to the desired size
    resized_images = F.interpolate(images, size=(size, size), mode='bilinear', align_corners=False)

    return resized_images

def Test_resize_function(num_images):
    toPIL = transforms.ToPILImage()
    train_loader = ggd.getdataloaders(num_images, 2)  # Set batch size to num_images
    for i, test_data in enumerate(train_loader):
        # Display the images in a grid
        fig, axs = plt.subplots(1, num_images, figsize=(20, 2))
        for j in range(num_images):
            image = test_data[j]
            #image = resize_images(image.unsqueeze(0),128)
            image = image.squeeze(0)
            axs[j].imshow(toPIL(image.cpu()))
            axs[j].axis('off')
        plt.show()
        print(test_data.shape)
        break

Test_resize_function(5)

import torch.optim as optim
from torch.autograd import Variable

def train_BigGAN2(train_loader, epochs, save_file_index):
  # Initialize the networks
  training_losses = []
  validations_losses = []
  training_accuracies = []
  validation_accuracies = []
  global image_saved_count
 
  # Training loop
  for epoch in range(epochs):
    loop = tqdm(total=len(train_loader), position=0, leave=False)
    for i, real_images in enumerate(train_loader):
      real_images = real_images.to(device)
      batch_size = real_images.size(0)
      # Update the discriminator
      optimizer_D.zero_grad()
      real_labels = Variable(torch.ones(batch_size, 1).to(device))
      fake_labels = Variable(torch.zeros(batch_size, 1).to(device))
      # Calculate the discriminator loss for real images
      real_output = discriminator(real_images)
      real_loss = wasserstein_loss(real_output, real_labels)
      # Generate fake images
      z = torch.randn(batch_size, 128).to(device)
      fake_images = generator(z)
      save_image = fake_images
      # Calculate the discriminator loss for fake images
      fake_output = discriminator(resize_images(fake_images.detach()))
      fake_loss = wasserstein_loss(fake_output, fake_labels)
           
      # Calculate the total discriminator loss and update the discriminator
      d_loss = torch.mean(fake_loss) - torch.mean(real_loss)
      d_loss.backward()
      optimizer_D.step()    
      # Update the generator
      optimizer_G.zero_grad()
           
        # Calculate the generator loss
      fake_output = discriminator(resize_images(fake_images))
      g_loss = wasserstein_loss(fake_output, real_labels)
           
        # Update the generator
      g_loss.backward()
      optimizer_G.step()

      loop.set_postfix(d_loss=d_loss.item(), g_loss=g_loss.item())
      loop.set_description(f"Epoch [{epoch + 1}/{epochs}]")  # Update epoch number
      loop.update(True)  # Refresh the progress bar
      if i % evolution_frequency == 0:
        save_images(save_image, evolution_folder, image_saved_count)
        image_saved_count += 1
   
  if SAVE_MODEL and epoch % 5 == 0:
    rcpf.save_checkpoint(generator, optimizer_G , save_file_index, filename= CHECKPOINT_GEN)
    rcpf.save_checkpoint(discriminator, optimizer_D, save_file_index, filename= CHECKPOINT_DISC)
    save_file_index = save_file_index + 1
  loop.close()
  return training_losses, training_accuracies

training_accuracies = []
training_losses = []
save_file_index = 0
try:
   training_losses, training_accuracies = train_BigGAN2(train_loader, epochs=10000, save_file_index = 1)
except Exception as e :
  print(e)
  torch.set_printoptions(profile = "default")
  import traceback
  print(traceback.format_exc())

rcpf.save_checkpoint(generator, optimizer_G , save_file_index, filename= CHECKPOINT_GEN)
rcpf.save_checkpoint(discriminator, optimizer_D, save_file_index, filename= CHECKPOINT_DISC)

# Plotting training losses
fig = plt.figure()
plt.title("Training Losses")
plt.xlabel("Iterations")
plt.ylabel("Loss")
plt.plot(training_losses, label='Training Loss', alpha=.5)
plt.legend()
plt.show()

# Plotting training and validation accuracies
fig = plt.figure()
plt.title("Training Accuracies")
plt.xlabel("Iterations")
plt.ylabel("Accuracy")
plt.plot(training_accuracies, label='Training Accuracy', alpha=.5)
plt.legend()
plt.show()

test_model(generator)

What did I do wrong ? did I Messed up the wasserstein_loss ?

def wasserstein_loss(output, target):
  return torch.mean(output * target)

why I am getting :
d_loss=-9.95e+3, g_loss=7.18e+3

I started training last night.

 

[WHT]

"torch.mean(output * target)"

Shouldn't a loss function be => output - target ?