training.py

# import Encoder_decoder_disc_VQ as Base_Models
import torch
from torch import nn
import cv2
import numpy as np
from torchvision import models
from torchsummary import summary
import matplotlib.pyplot as plt # plotting library
import numpy as np # this module is useful to work with numerical arrays
#import pandas as pd 
import random 
import torch.utils.data as data_loader
import torchvision
from torchvision import transforms
import torch.nn.functional as F
import Encoder_decoder_disc_VQ as Base_Models


if torch.cuda.is_available():
    device="cuda" 
    torch.cuda.empty_cache()

else:
    device="cpu"

print("the device is : ",device)


Disc=Base_Models.Disc_net1().to(device)

encode=Base_Models.Encoder().to(device)
VQ_=Base_Models.VQ().to(device)
decode=Base_Models.Decoder().to(device)

print(summary(encode, (3, 128, 128)))

print(summary(decode, ( 128, 12, 12)))
print(summary(Disc, ( 3, 128, 128)))


dataset = torchvision.datasets.ImageFolder(root="dataset224\\train",
                           transform=transforms.Compose([
                               transforms.Resize((128,128)),
                               transforms.ToTensor()
                           ]))
# Create the dataloader
dataloader = torch.utils.data.DataLoader(dataset, batch_size=128,
                                         shuffle=True, num_workers=2)
#dataloader=data_loading("C")
criterion = nn.BCELoss()
# criterion = nn.CrossEntropyLoss()
rec_loss = nn.MSELoss()



lr=0.0001
lr=0.0001
params_to_optimize = [{'params': VQ_.parameters()},
    {'params': encode.parameters()},
    {'params': decode.parameters()}
]


optimizer_gen = torch.optim.Adam(params_to_optimize, lr=0.0001)

#optimizerGen = torch.optim.Adam(params_to_optimize, lr=lr)
optimizer_disc = torch.optim.Adam(Disc.parameters(), lr=0.0001)



real_label = 1.
fake_label = 0.


def show_image1(img):
    npimg = img.numpy()
    npimg=(np.transpose(npimg, (1, 2, 0)))*255
    npimg = cv2.cvtColor(npimg, cv2.COLOR_RGB2BGR)
    cv2.imwrite("test10.jpg",npimg)
    #plt.imshow(np.transpose(npimg, (1, 2, 0)))
def train_epoch(encode, decode, disc, vq, device, dataloader, criterion, rec_loss_fn, optimizer_gen, optimizer_disc):
    encode.train()
    decode.train()
    disc.train()
    vq.train()

    running_loss = 0.0
    running_disc_loss = 0.0
    running_gen_loss = 0.0
    for images, _ in dataloader:
        images = images.to(device)
        b_size = images.size(0)

        # Discriminator update
        optimizer_disc.zero_grad()

        # Train with real images
        real_labels = torch.full((b_size,), 1., dtype=torch.float, device=device)
        
        output_real = disc(images).view(-1)
        # print(output_real.shape,real_labels.shape)
        loss_real = criterion(output_real, real_labels)
        loss_real.backward()

        # Train with fake images
        with torch.no_grad():
            encoded_images = encode(images)
            _, quantized_images, _, _ = vq(encoded_images)
            fake_images = decode(quantized_images)
        fake_labels = torch.full((b_size,), 0., dtype=torch.float, device=device)
        output_fake = disc(fake_images.detach()).view(-1)
        loss_fake = criterion(output_fake, fake_labels)
        loss_fake.backward()

        # Update discriminator
        optimizer_disc.step()
        disc_loss = loss_real.item() + loss_fake.item()

        # Generator and VQ-VAE update
        optimizer_gen.zero_grad()

        # Forward pass through VQ-VAE
        encoded_images = encode(images)
        q_loss, quantized_images,perplexity,encodings = vq(encoded_images)
        decoded_images = decode(quantized_images)

        # Reconstruction loss
        rec_loss = rec_loss_fn(decoded_images, images)

        # Fooling the discriminator loss
        output = disc(decoded_images).view(-1)
        gen_loss = criterion(output, real_labels)

        # Combined loss
        loss = rec_loss + gen_loss + q_loss
        loss.backward()

        # Update generator and VQ-VAE
        optimizer_gen.step()

        running_loss += loss.item()
        running_disc_loss += disc_loss
        running_gen_loss += gen_loss.item()
        img_recon = decoded_images.cpu()
        images = images.cpu()
        img_recon=torch.cat((img_recon, images), 0)
        codebook = encodings.weight.cpu()
        show_image1(torchvision.utils.make_grid(img_recon[:],16,20))
        codebook = codebook.detach().numpy()
        cv2.imwrite("code.jpg",(codebook*255))

    avg_loss = running_loss / len(dataloader)
    avg_disc_loss = running_disc_loss / len(dataloader)
    avg_gen_loss = running_gen_loss / len(dataloader)
    print(f'Train loss: {avg_loss:.4f}, Disc loss: {avg_disc_loss:.4f}, Gen loss: {avg_gen_loss:.4f}')

    return avg_loss, avg_disc_loss, avg_gen_loss,encode,decode,disc,vq






num_epochs =400
diz_loss = {'generator_loss':[]}
diz2_loss = {'discr_loss':[]}


import time
num_epochs = 1000
for epoch in range(num_epochs):
    start_time = time.time()
    avg_loss, avg_disc_loss, avg_gen_loss,encode,decode,disc,vq = train_epoch(
    encode=encode,
    decode=decode,
    disc=Disc,
    vq=VQ_,
    device=device,
    dataloader=dataloader,
    criterion=criterion,
    rec_loss_fn=rec_loss,
    optimizer_gen=optimizer_gen,
    optimizer_disc=optimizer_disc
)

    print(f'\n Epoch {epoch+1}/{num_epochs} \t Gen loss {avg_gen_loss} \t Disc loss {avg_disc_loss}')
    print("Time of epoch:", time.time() - start_time)
    diz_loss['generator_loss'].append(avg_gen_loss)
    diz2_loss['discr_loss'].append(avg_disc_loss)

# Visualization of Losses
plt.figure(figsize=(10, 8))
plt.semilogy(diz_loss['generator_loss'], label='Generator Loss')
plt.semilogy(diz2_loss['discr_loss'], label='Discriminator Loss')
plt.xlabel('Epoch')
plt.ylabel('Average Loss')
plt.legend()
plt.show()