SS_disc_task.py

import os
import torch
import torch.nn as nn
from torch.autograd import Variable
import numpy as np
import random
from matplotlib import pyplot as plt
from tqdm import tqdm
from torch import optim
import argparse
import sys
import pandas as pd

from source.latent_module import Rep_AE, LatentCond_Unet, RepDDPM
from source.DDPM import Diffusion
from source.SemiSup_module import top_model, SemiSup_DDPM
import source.helper_func as hf
import source.preprocess as prep

import logging
from torch.utils.tensorboard import SummaryWriter
import torchvision.datasets as datasets
from torch.utils.data import DataLoader, Dataset
from torchvision import transforms as T
import math




def set_seed(SEED: int):

    torch.manual_seed(SEED)
    random.seed(SEED)
    np.random.seed(SEED)

def prepare_data(
        batch_size: int = 32,
        perc_size: float = 0.01, 
    ):

    transforms = T.Compose(
            [
                T.Resize([32,32]),
                T.ToTensor(),
                T.Normalize((0.5), (0.5))
            ]
    )

    mnist_trainset = datasets.MNIST(root='./data', train=True, download=True, transform=transforms)
    mnist_testset = datasets.MNIST(root='./data', train=False, download=True, transform=transforms)
    
    #train_dataloader = torch.utils.data.DataLoader(mnist_trainset, batch_size=batch_size)
    dataset_size = mnist_trainset.train_labels.shape[0]
    y_accept = torch.ones((dataset_size))
    y_accept[int(perc_size * dataset_size):] *= 0

    trainset = prep.SS_MNIST_dataset(
            x_onehot = mnist_trainset.train_data,
            y_class = mnist_trainset.train_labels,
            L_disc_accept = y_accept
    )

    train_dataloader = DataLoader(trainset, batch_size = batch_size, num_workers = 4, shuffle = True)

    test_dataloader = torch.utils.data.DataLoader(mnist_testset, batch_size=batch_size)


    return train_dataloader, test_dataloader


def get_args():

    parser = argparse.ArgumentParser()

    parser.add_argument('-S', '--SEED', default = 42, help = 'flag: random seed', type = int)
    parser.add_argument('-bs', '--batch_size', default = 32, help = 'flag: batch size', type = int)
    parser.add_argument('-lr', '--learning_rate', default = 1e-3, help = 'flag: learning rate', type = float)
    parser.add_argument('-e', '--epochs', default = 10, help = 'flag: training epochs', type = int)
    parser.add_argument('-ps', '--perc_size', default = '0.01', help = 'flag: training set size for disc loss', type = str)

    # model hyperparameter
    parser.add_argument('-ig', '--img_size', default = 32, help = 'flag: image size', type = int)
    parser.add_argument('-cin', '--c_in', default = 1, help = 'flag: input channel to the Unet model', type = int)
    parser.add_argument('-cout', '--c_out', default = 1, help = 'flag: output channel to the Unet model', type = int)
    parser.add_argument('-fnc', '--first_num_channel', default = 64, help = 'flag: number of conv channels for the first layer', type = int)
    parser.add_argument('-td', '--time_dim', default = 256, help = 'flag: embedding dimension for the time positions', type = int)
    parser.add_argument('-nl', '--num_layers', default = 3, help = 'flag: number of layers', type = int)
    parser.add_argument('-bn', '--bn_layers', default = 2, help = 'flag: number of Unets bottleneck layers', type = int)
    parser.add_argument('-rl', '--rep_learning', default = False, help = 'flag: choose to use latent conditional info or not.', type = bool)
    parser.add_argument('-zd', '--z_dim', default = 5, help = 'flag: latent space dimension', type = int)
    parser.add_argument('-dl', '--disc_num_layers', default = 2, help = 'flag: classifier depth', type = int)
    parser.add_argument('-dw', '--disc_width', default = 1000, help = 'flag: classifier width', type = int)



    # learning prefactors
    parser.add_argument('-al', '--alpha', default = 0.99, help = 'flag: mutual information (x,z) prefactor weight', type = float)
    parser.add_argument('-be', '--beta', default = 10, help = 'flag: MMD prefactor weight', type = float)
    parser.add_argument('-ga', '--gamma', default = 10, help = 'flag: discriminative prefactor weight', type = float)
    parser.add_argument('-x', '--xi', default = 1, help = 'flag: NLL prefactor weight', type = float)
       

    # filepaths
    parser.add_argument('-fop', '--fig_output_path', default = './output/SSlatentDDPM/mnist_samples_', help = 'flag: filenames for the mnist files', type = str)
    parser.add_argument('-omp', '--output_model_path', default = './output/SSlatentDDPM/model_pretrained_weights.pth', help = 'flag: model weights', type = str)
    parser.add_argument('-ocp', '--output_csv_path', default = './output/SSlatentDDPM/model_train_history.csv', help = 'flag: model training/testing loss history', type = str)
    parser.add_argument('-tfo', '--tsne_path', default = './output/SSlatentDDPM/tsne_plot', help = 'flag: ', type = str)

    args = parser.parse_args()

    return args


def compute_loss(
        model,
        true_noise,
        pred_noise,
        z_true,
        z_infer,
        z_mu,
        z_var,
        y_true,
        y_pred,
        y_accept
    ):

    # DDPM loss
    mse = nn.MSELoss(reduction = 'none')
    loss_ddpm = torch.mean(
                        torch.sum(
                                mse(true_noise, pred_noise), dim = (-2, -1)
                        )
    )

    # KL divergence
    loss_kl = model.latent_AE.kl_loss(z_mu, z_var)

    # max mean discrepancy loss
    loss_mmd = model.latent_AE.mmd_loss(z_true, z_infer)

    # discriminative loss:
    ce_loss = nn.CrossEntropyLoss(reduction = 'none')
    y_true = torch.argmax(y_true, dim = -1).long()
    loss_disc = ce_loss(y_pred, y_true)
    loss_disc_accept = loss_disc * y_accept # note: we multiple zeros to batch indexes that are not included
    loss_disc_accept = torch.mean(loss_disc_accept, dim = -1)
    loss = model.xi * loss_ddpm + (1-model.alpha)*loss_kl + (model.alpha + model.beta - 1)*loss_mmd + model.gamma*loss_disc_accept

    return loss, loss_ddpm, loss_kl, loss_mmd, loss_disc_accept

def compute_acc(true, pred):
    return sum(np.array(true) == np.array(pred)) / len(true)

@torch.no_grad()
def end_epoch_validation(
        DEVICE: str,
        model: any,
        diffusion: any,
        test_dataloader: any
    ):
    
    model.eval()
    onehot_encoding = torch.eye(10)
    
    y_test_preds, y_test_true = [], []

    epoch_loss, epoch_mse_loss, epoch_kld_loss, epoch_mmd_loss, epoch_disc_loss = [],  [], [], [], []
    for iteration, valid_batch in enumerate(test_dataloader):
        
        # validation batch using the testing set
        valid_X0 = Variable(valid_batch[0].to(DEVICE))
        valid_y = Variable(onehot_encoding[
            valid_batch[1]].to(DEVICE)
        )
        
        # timesteps for each batch
        t = diffusion.sample_timesteps(valid_X0.shape[0]).to(DEVICE)    
            
        # corrupt training batch
        x_t, noise = diffusion.noise_images(valid_X0, t)

        # predict noise 
        predicted_noise, z, z_mean, z_var, y_pred = model(x_t, t, valid_X0)
        
        # sample from some noise:

        # normal dist:
        if DEVICE == 'cuda':
            z_true_samples = torch.randn_like(z).to(DEVICE)
        
        else:
            z_true_samples = torch.randn_like(z)
        
    
        # loss
        loss, loss_ddpm, loss_kld, loss_mmd, loss_disc = model.compute_loss(
                                        true_noise = noise,
                                        pred_noise = predicted_noise,
                                        z_true = z_true_samples,
                                        z_infer = z,
                                        z_mu = z_mean,
                                        z_var = z_var,
                                        y_true = valid_y,
                                        y_pred = y_pred,
        )

        epoch_loss.append(loss.item())
        epoch_mse_loss.append(loss_ddpm.item())
        epoch_kld_loss.append(loss_kld.item())
        epoch_mmd_loss.append(loss_mmd.item())
        epoch_disc_loss.append(loss_disc.item())
        
        # save predictions
        y_test_true += list(torch.argmax(valid_y, dim = -1).cpu().numpy())
        y_test_preds += list(torch.argmax(y_pred, dim = -1).cpu().numpy())

    model.train()
    
    print(f'Validation | loss: {np.mean(epoch_loss)}'
        + f'| Mean squared error: {np.mean(epoch_mse_loss)}'
        + f'| KLD loss: {np.mean(epoch_kld_loss)}'
        + f'| MMD loss: {np.mean(epoch_mmd_loss)}'
        + f'| Disc loss: {np.mean(epoch_disc_loss)}'
    )
    
    y_acc = compute_acc(y_test_true, y_test_preds)
    
    print(f'Validation accuracy:', y_acc)


    return np.mean(epoch_loss), np.mean(epoch_mse_loss), np.mean(epoch_kld_loss), np.mean(epoch_mmd_loss), np.mean(epoch_disc_loss), y_acc
       

def train_model(
    train_dataloader: any,
    test_dataloader: any,
    args: any,
    DEVICE: str = 'cuda'
    ) -> (any, dict):
    
    # model hyperparameters
    img_size = args.img_size
    c_in = args.c_in
    c_out = args.c_out
    first_num_channel = args.first_num_channel
    time_dim = args.time_dim
    num_layers = args.num_layers
    bn_layers = args.bn_layers
    z_dim = args.z_dim
    disc_num_layers = args.disc_num_layers
    disc_width = args.disc_width

    # loss prefactors
    xi = args.xi
    alpha = args.alpha
    beta = args.beta
    gamma = args.gamma

    # configure model
    rep_AE = Rep_AE(
            img_size = img_size,
            c_in = c_in,
            c_out = c_out,
            first_num_channel = first_num_channel,
            z_dim = z_dim,
            time_dim = time_dim,
            num_layers = num_layers,
            bn_layers = bn_layers
    ).to(DEVICE)

    tm_disc = top_model(
            num_layers = disc_num_layers,
            hidden_width = disc_width,
            z_dim = z_dim,
            num_classes = 10
    ).to(DEVICE)

    Unet = LatentCond_Unet(
            c_in = c_in + c_in,
            c_out = c_out,
            first_num_channel = first_num_channel,
            time_dim = time_dim,
            num_layers = num_layers,
            bn_layers = bn_layers
    ).to(DEVICE)

    
    diffusion = Diffusion(
            img_size = img_size,
            device =  DEVICE,
            schedule = 'linear'
    )

    model = SemiSup_DDPM(
            Unet = Unet,
            latent_AE = rep_AE,
            top_model = tm_disc,
            diffusion = diffusion,
            xi = xi,
            alpha = alpha,
            beta = beta,
            gamma = gamma
    ).to(DEVICE)


    optimizer = optim.AdamW(model.parameters(), lr = args.learning_rate)
    mse = nn.MSELoss(reduction = 'none')
    model.train()
    
    history_dict = {
        'epoch': list(),
        'train_loss': list(),
        'test_loss': list(),
        'train_mse_loss': list(),
        'test_mse_loss': list(),
        'train_kld_loss': list(),
        'test_kld_loss': list(),
        'train_mmd_loss': list(),
        'test_mmd_loss': list(),
        'train_disc_loss': list(),
        'test_disc_loss': list(),
        'test_acc': list(),
    }
    
    onehot_encoding = torch.eye(10)
    
    # train model
    for epoch in range(args.epochs):
        
        epoch_loss, epoch_mse_loss, epoch_kld_loss, epoch_mmd_loss, epoch_disc_loss = [], [], [], [], []
        for iteration, train_batch in tqdm(enumerate(train_dataloader)):
            
            # training batch
            train_X0 = Variable(
                            train_batch[0].to(DEVICE)
            )
            train_y = Variable(
                            onehot_encoding[train_batch[1].long()].to(DEVICE)
            )
            train_y_accept = Variable(
                            train_batch[2].to(DEVICE)
            )

            # timesteps for each batch
            t = diffusion.sample_timesteps(train_X0.shape[0]).to(DEVICE)

            # corrupt training batch
            x_t, noise = diffusion.noise_images(train_X0, t)

            # predicted noise
            predicted_noise, z, z_mean, z_var, y_pred = model(x = x_t, t = t, x0 = train_X0)
            
            # sample from some noise
            if DEVICE == 'cuda':
                z_true_samples = Variable(torch.randn_like(z)).to(DEVICE)
            else:
                z_true_samples = torch.randn_like(z)

            # loss
            loss, loss_ddpm, loss_kld, loss_mmd, loss_disc = compute_loss(
                                                    model = model,
                                                    true_noise = noise,
                                                    pred_noise = predicted_noise,
                                                    z_true = z_true_samples,
                                                    z_infer = z,
                                                    z_mu = z_mean,
                                                    z_var = z_var,
                                                    y_true = train_y,
                                                    y_pred = y_pred,
                                                    y_accept = train_y_accept
            )

            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            epoch_loss.append(loss.item())
            epoch_mse_loss.append(loss_ddpm.item())
            epoch_kld_loss.append(loss_kld.item())
            epoch_mmd_loss.append(loss_mmd.item())
            epoch_disc_loss.append(loss_disc.item())
            

        train_loss = np.mean(epoch_loss)
        train_mse_loss = np.mean(epoch_mse_loss)
        train_kld_loss = np.mean(epoch_kld_loss)
        train_mmd_loss = np.mean(epoch_mmd_loss)
        train_disc_loss = np.mean(epoch_disc_loss)

        print(f'Training | loss: {train_loss} '
            + f'| Mean squared error: {train_mse_loss} ' 
            + f'| KLD loss: {train_kld_loss} ' 
            + f'| MMD loss: {train_mmd_loss}'
            + f'| disc loss: {train_disc_loss}'
        )
        
        history_dict['epoch'].append(epoch)
        history_dict['train_loss'].append(train_loss)
        history_dict['train_mse_loss'].append(train_mse_loss)
        history_dict['train_kld_loss'].append(train_kld_loss)
        history_dict['train_mmd_loss'].append(train_mmd_loss)
        history_dict['train_disc_loss'].append(train_disc_loss)
        
        test_loss, test_mse_loss, test_kld_loss, test_mmd_loss, test_disc_loss, test_acc = end_epoch_validation(
            DEVICE = DEVICE,
            model = model,
            diffusion = diffusion,
            test_dataloader = test_dataloader
        )
    
        history_dict['test_loss'].append(test_loss)
        history_dict['test_mse_loss'].append(test_mse_loss)
        history_dict['test_kld_loss'].append(test_kld_loss)
        history_dict['test_mmd_loss'].append(test_mmd_loss)
        history_dict['test_disc_loss'].append(test_disc_loss)
        history_dict['test_acc'].append(test_acc)

        # generate samples to quantify quality:
        sample_images = hf.sample(diffusion, model, 20, DEVICE)
        hf.plot_random_samples(sample_images)
        plt.savefig(f'{args.fig_output_path}[epoch={epoch}].png', dpi = 300)
 

    return model, history_dict


if __name__ == '__main__':

    # get arguments
    args = get_args()
    
    SEED = args.SEED # random seed
    batch_size = args.batch_size # batch size
    output_model_path = args.output_model_path # model output path
    output_csv_path = args.output_csv_path # train/test loss history path
    perc_sizes = [float(item) for item in args.perc_size.split(',')] # percentage size

    # set seed for reproducibility...
    set_seed(SEED = SEED)

    # check GPU
    if torch.cuda.is_available():
        print('GPU available')

    else:
        print('Please enable GPU and rerun script')
        quit()


    USE_CUDA = True
    DEVICE = 'cuda' if USE_CUDA else 'cpu'

    SS_hist_dict = {

    }

    for perc_size in perc_sizes:

        # prepare data
        train_dataloader, test_dataloader = prepare_data(
                                                batch_size=batch_size,
                                                perc_size = perc_size
        )


        # train model
        model, hist_dict = train_model(
            train_dataloader = train_dataloader,
            test_dataloader = test_dataloader,
            args = args,
            DEVICE = DEVICE
        
        )

        SS_hist_dict[f'train_loss[perc_size={perc_size}]'] = hist_dict['train_loss']
        SS_hist_dict[f'test_loss[perc_size={perc_size}]'] = hist_dict['test_loss']
        SS_hist_dict[f'train_mse_loss[perc_size={perc_size}]'] = hist_dict['train_mse_loss']
        SS_hist_dict[f'test_mse_loss[perc_size={perc_size}]'] = hist_dict['test_mse_loss']
        SS_hist_dict[f'train_kld_loss[perc_size={perc_size}]'] = hist_dict['train_kld_loss']
        SS_hist_dict[f'test_kld_loss[perc_size={perc_size}]'] = hist_dict['test_kld_loss']
        SS_hist_dict[f'train_mmd_loss[perc_size={perc_size}]'] = hist_dict['train_mmd_loss']
        SS_hist_dict[f'test_mmd_loss[perc_size={perc_size}]'] = hist_dict['test_mmd_loss']
        SS_hist_dict[f'test_acc[perc_size={perc_size}]'] = hist_dict['test_acc']



        torch.save(model.state_dict(), f'{output_model_path}[perc_size={perc_size}].pth')
    
    
    # dataframe of the history dict
    SS_hist_df = pd.DataFrame(SS_hist_dict)

    # save dataframe
    SS_hist_df.to_csv(f'{output_csv_path}', index = False)