hp_latentDDPM.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
import logging
from torch.utils.tensorboard import SummaryWriter
import torchvision.datasets as datasets
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
    ):

    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)
    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)
    
    # 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('-zi', '--z_init', default = 1, help = 'flag: init latent space dimension', type = int)
    parser.add_argument('-zf', '--z_final', default = 10, help = 'flag: final latent space dimension', 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)
    

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

    args = parser.parse_args()

    return args

@torch.no_grad()
def sample(
        diffusion: any,
        model: any,
        n: int = 20
    ):
    
    model.eval()
    
        
    x_sample = torch.randn((n, 1, diffusion.img_size, diffusion.img_size)).to(diffusion.device)
    z_0 = torch.randn((n, model.latent_AE.z_dim)).to(DEVICE)

    for ii in tqdm(reversed(range(1, diffusion.noise_steps)), position = 0):
            
        t = (torch.ones(n) * ii).long().to(diffusion.device)
            
        # decode latents
        x_z = model.sample_randomly(z_0, t)

        # predict noise
        predicted_noise = model.Unet(x_sample, x_z, t)

        alpha = diffusion.alpha[t][:, None, None, None]
        alpha_hat = diffusion.alpha_hat[t][:, None, None, None]
        beta = diffusion.beta[t][:, None, None, None]
            
        if ii > 1:
            noise = torch.randn_like(x_sample)
        else:
            noise = torch.zeros_like(x_sample)
                
                
        alpha_prefactor = (1-alpha)/ (torch.sqrt(1-alpha_hat))

        x_sample = 1 / torch.sqrt(alpha) * (x_sample - alpha_prefactor*predicted_noise) + torch.sqrt(beta) * noise
            
    model.train()
    x_sample = (x_sample.clamp(-1, 1) + 1)/ 2
    return x_sample


def plot_samples(
        fig_output_path: str,
        X_sample: torch.FloatTensor,
        epoch: int
    ):

    fig, axes = plt.subplots(5,4, dpi = 300, sharex = True, sharey = True)

    fig.subplots_adjust(hspace = 0, wspace = 0)

    idx = 0
    for ii in range(5):

        for jj in range(4):

            axes[ii,jj].imshow(
                    X_sample[idx, 0].cpu()
            )
            axes[ii,jj].axis('off')

            axes[ii,jj].set_xticklabels([])
            axes[ii,jj].set_yticklabels([])
            idx+=1
    plt.tight_layout()
    plt.savefig(f'{fig_output_path}_epoch={epoch}.png', dpi = 300)
    

@torch.no_grad()
def end_epoch_validation(
        DEVICE: str,
        model: any,
        diffusion: any,
        test_dataloader: any
    ):
    
    model.eval()
    
    epoch_loss, epoch_mse_loss, epoch_kld_loss, epoch_mmd_loss = [], [], [], []
    for iteration, valid_batch in enumerate(test_dataloader):
        
        # validation batch using the testing set
        valid_batch = Variable(valid_batch[0].to(DEVICE))
        
        # timesteps for each batch
        t = diffusion.sample_timesteps(valid_batch.shape[0]).to(DEVICE)    
            
        # corrupt training batch
        x_t, noise = diffusion.noise_images(valid_batch, t)

        # predict noise 
        predicted_noise, z, z_mean, z_var = model(x_t, t, valid_batch)
        
        # 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 = 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,
                                        gamma = 1
        )

        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())
        
    model.train()
    
    print(f'Validation | loss: {np.mean(epoch_loss)} | Mean squared error: {np.mean(epoch_mse_loss)} | KLD loss: {np.mean(epoch_kld_loss)} | MMD loss: {np.mean(epoch_mmd_loss)}')

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

def train_model(
    z_dim: int,
    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

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

    # 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)

    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 = RepDDPM(
            Unet = Unet,
            latent_AE = rep_AE,
            diffusion = diffusion,
            alpha = alpha,
            beta = beta
    ).to(DEVICE)


    optimizer = optim.Adam(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 model
    for epoch in range(args.epochs):
        
        epoch_loss, epoch_mse_loss, epoch_kld_loss, epoch_mmd_loss = [], [], [], []
        for iteration, train_batch in tqdm(enumerate(train_dataloader)):

            # training batch
            train_batch = Variable(train_batch[0].to(DEVICE))

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

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

            # predicted noise
            predicted_noise, z, z_mean, z_var = model(x = x_t, t = t, x0 = train_batch)
            
            # 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 = 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,
                                                    gamma = 1
            )

            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())
            
        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)

        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}')
        
        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)
        
        test_loss, test_mse_loss, test_kld_loss, test_mmd_loss = 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)
        

    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
    z_init = args.z_init # initial latent space hp
    z_final = args.z_final # final latent space hp

    # 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'


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

    hp_hist_dict = {
            }
    for z_dim in range(z_init, z_final):
        # train model
        model, hist_dict = train_model(
            z_dim = z_dim,
            train_dataloader = train_dataloader,
            test_dataloader = test_dataloader,
            args = args,
            DEVICE = DEVICE
        )     
    
        
        hp_hist_dict[f'train_loss[z_dim={z_dim}]'] = hist_dict['train_loss']
        hp_hist_dict[f'test_loss[z_dim={z_dim}]'] = hist_dict['test_loss']
        hp_hist_dict[f'train_mse_loss[z_dim={z_dim}]'] = hist_dict['train_mse_loss']
        hp_hist_dict[f'test_mse_loss[z_dim={z_dim}]'] = hist_dict['test_mse_loss']
        hp_hist_dict[f'train_kld_loss[z_dim={z_dim}]'] = hist_dict['train_kld_loss']
        hp_hist_dict[f'test_kld_loss[z_dim={z_dim}]'] = hist_dict['test_kld_loss']
        hp_hist_dict[f'train_mmd_loss[z_dim={z_dim}]'] = hist_dict['train_mmd_loss']
        hp_hist_dict[f'test_mmd_loss[z_dim={z_dim}]'] = hist_dict['test_mmd_loss']
            

        # save model
        torch.save(model.state_dict(), f'{output_model_path}[z_dim={z_dim}].pth')


    # dataframe of the history dict
    hp_hist_df = pd.DataFrame(hp_hist_dict)

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