time_prediction_training.py

import json
import argparse
from core.logger import load_json
from data.time_predictor_dataset import TimePredictorDataset
from core.experiment_directory_setup import get_workdir
from data.split_dataset import DataLocation
from model.ddpm_modules.time_predictor import TimePredictor
from torch.utils.data import DataLoader
import torch.optim as optim
import torch
from torch.optim import Adam
from core.wandb_logger import WandbLogger
from tqdm import tqdm
import numpy as np
from core.logger import mkdirs
import os
from split import add_git_info
from model.normalizer import NormalizerXT
from predtiler.dataset import get_tiling_dataset, get_tile_manager

def get_normalizer(dset, opt, num_pixels=500_000 * 128 * 128, return_data_arr=False):
    
    original_flag = dset.set_random_patching(True)
    xt_normalizer = NormalizerXT()
    if 'normalize_channels' in opt['datasets'] and opt['datasets']['normalize_channels'] is True:
        pass
    else:
        dset.reset_fixed_t()
        data_mean_arr = []
        data_std_arr = []
        cnt = 0
        idx = 0
        for _ in tqdm(range(10000)):
            val_loader = torch.utils.data.DataLoader(dset, batch_size=16, shuffle=True, num_workers=0)
            bar = tqdm(val_loader)
            for (x, t_float) in bar:
                idx+=1
                x = x.cuda()
                t_float = t_float.cuda()
                x = xt_normalizer.normalize(x,t_float, update=True)
                cnt += np.prod(x.shape)
                bar.set_description(f'{cnt//10e6}M/{num_pixels//10e6}M pixels processed')
                if idx %10 == 0 and return_data_arr:
                    data_mean_arr.append(xt_normalizer.data_mean.cpu().numpy())
                    data_std_arr.append(xt_normalizer.data_std.cpu().numpy())
                if cnt > num_pixels:
                    break
            if cnt > num_pixels:
                break

        # plt.plot(np.stack(data_mean_arr)[:,90])
    dset.set_random_patching(original_flag)
    if return_data_arr:
        return xt_normalizer, data_mean_arr, data_std_arr
    return xt_normalizer


def get_datasets(opt, tiled_pred=False):
    patch_size = opt['datasets']['patch_size']
    target_channel_idx = opt['datasets'].get('target_channel_idx', None)
    upper_clip = opt['datasets'].get('upper_clip', None)
    max_qval = opt['datasets']['max_qval']
    channel_weights = opt['datasets'].get('channel_weights', None)
    normalize_channels = opt['datasets'].get('normalize_channels', False)
    data_type = opt['datasets']['train']['name']  
    uncorrelated_channels = opt['datasets']['train']['uncorrelated_channels']
    
    assert data_type in ['cifar10', 'Hagen','COSEM_jrc-hela', 'HT_LIF24', "BioSR", "HT_T24", "PaviaATN",
                         'COSEM_jrc-choroid-plexus-2', 'goPro2017dehazing'], f'Invalid data type: {data_type}'
    if data_type == 'Hagen':
        train_data_location = DataLocation(channelwise_fpath=(opt['datasets']['train']['datapath']['ch0'],
                                                        opt['datasets']['train']['datapath']['ch1']))
        val_data_location = DataLocation(channelwise_fpath=(opt['datasets']['val']['datapath']['ch0'],
                                                        opt['datasets']['val']['datapath']['ch1']))
    elif data_type in ['cifar10', 'HT_LIF24', 'COSEM_jrc-hela', "BioSR", "HT_T24",'COSEM_jrc-choroid-plexus-2','goPro2017dehazing', "PaviaATN"]:
        train_data_location = DataLocation(directory=(opt['datasets']['train']['datapath']))
        val_data_location = DataLocation(directory=(opt['datasets']['val']['datapath']))
        if data_type == 'goPro2017dehazing':
            train_data_location.datasplit_type = 'train'
            val_data_location.datasplit_type = 'val'
            train_data_location.limit_count = opt['datasets']['train'].get('limit_count', None)
            val_data_location.limit_count = opt['datasets']['val'].get('limit_count', None)
    else:
        raise ValueError('Invalid data type')
    
    gaussian_noise_std_factor = opt['datasets']['train'].get('gaussian_noise_std_factor', None)

    train_set = TimePredictorDataset(data_type, train_data_location, patch_size, 
                             target_channel_idx=target_channel_idx, 
                                max_qval=max_qval, upper_clip=upper_clip,
                                uncorrelated_channels=uncorrelated_channels,
                                channel_weights=channel_weights,
                             normalization_dict=None, enable_transforms=True,random_patching=True,
                             gaussian_noise_std_factor=gaussian_noise_std_factor,
                             normalize_channels=normalize_channels)

    if not tiled_pred:
        class_obj = TimePredictorDataset 
    else:
        if data_type == 'Hagen':
            data_shape = (10, 2048, 2048)
        elif data_type in ['HT_LIF', 'HT_LIF24']:
            data_shape = (10, 1608, 1608)
        elif data_type == 'COSEM_jrc-hela':
            data_shape = (96, 900, 1400)
        elif data_type == 'COSEM_jrc-choroid-plexus-2':
            data_shape = (96, 900, 1220)
        elif data_type =='HT_T24':
            data_shape = (36, 1608, 1608)
        elif data_type == 'BioSR':
            data_shape = (5, 1004, 1004)
        else:
            raise ValueError('Invalid data type')
        
        tile_manager = get_tile_manager(data_shape, (1, 64, 64), (1, patch_size, patch_size))

        class_obj = get_tiling_dataset(TimePredictorDataset, tile_manager)
        # raise NotImplementedError('Tiled prediction not implemented yet')

    val_set = class_obj(data_type, val_data_location, patch_size, target_channel_idx=target_channel_idx,
                           normalization_dict=train_set.get_input_target_normalization_dict(),
                           max_qval=max_qval,
                            upper_clip=upper_clip,
                            channel_weights=channel_weights,
                           enable_transforms=False,
                            random_patching=False,
                            normalize_channels=normalize_channels)
    return train_set, val_set

def start_training(opt):
    if opt['enable_wandb']:
        import wandb
        add_git_info(opt)
        wandb_logger = WandbLogger(opt, opt['path']['experiment_root'], opt['experiment_name'])
        # wandb.define_metric('validation/val_step')
        # wandb.define_metric('epoch')
        # wandb.define_metric("validation/*", step_metric="val_step")
        # val_step = 0
    else:
        wandb_logger = None

    train_set, val_set = get_datasets(opt, tiled_pred=False)
    model_opt = opt['model']
    model_kwargs = {}
    model_kwargs['scale_augmentation'] = model_opt.get('scale_augmentation', False)
    
    if model_kwargs['scale_augmentation']:
        model_kwargs['scale_augmentation_delta'] = model_opt['scale_augmentation_delta']
    
    model = TimePredictor(
        in_channel=model_opt['unet']['in_channel'],
        out_channel=model_opt['unet']['out_channel'],
        norm_groups=model_opt['unet']['norm_groups'],
        inner_channel=model_opt['unet']['inner_channel'],
        channel_mults=model_opt['unet']['channel_multiplier'],
        attn_res=model_opt['unet']['attn_res'],
        res_blocks=model_opt['unet']['res_blocks'],
        dropout=model_opt['unet']['dropout'],
        image_size=opt['datasets']['patch_size'],
        initial_instance_norm= model_opt['unet'].get('initial_instance_norm', False),
        **model_kwargs,
        )
    model = model.cuda()

    dummy_normalizer_flag = opt['datasets'].get('normalize_channels', False) is True
    # instantiate the normalizer
    if dummy_normalizer_flag:
        print('--------Dummy Normalizer Activated--------')
        xt_normalizer_train = None
        xt_normalizer_val = None
    else:
        xt_normalizer_train = get_normalizer(train_set, opt)
        # we donot need a separate normalizer for validation set. 
        # this is because training and validation set are drawn from the same distribution.
        xt_normalizer_val = get_normalizer(val_set, opt)

    train_loader = DataLoader(train_set, batch_size=opt['datasets']['train']['batch_size'], shuffle=True, num_workers=opt['datasets']['train']['num_workers'])
    val_loader = DataLoader(val_set, batch_size=opt['datasets']['train']['batch_size'], shuffle=False, num_workers=opt['datasets']['train']['num_workers'])

    optimizer = Adam(model.parameters(), lr=opt['train']['optimizer']['lr'])
    lr_scheduler_patience = opt['train']['lr_scheduler_patience']
    # learning rate scheduler.
    scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer,
                                                            mode='min',
                                                            patience=lr_scheduler_patience,
                                                            factor=0.5,
                                                            min_lr=1e-6,
                                                            verbose=True)# create loss function.
    if opt['model']['loss_type'] == 'l1':
        loss_fn = torch.nn.L1Loss()
    elif opt['model']['loss_type'] == 'l2':
        loss_fn = torch.nn.MSELoss()
    else:
        raise ValueError(f"Invalid loss type {opt['model']['loss_type']}")   

    # tqdm bar with loss 
    epoch_losses = []
    num_epochs = opt['train']['num_epochs']
    best_val_loss = 1e6
    for epoch in range(num_epochs):
        bar = tqdm(enumerate(train_loader))
        loss_arr = []
        for i, (x, t_float) in bar:
            model.train()
            optimizer.zero_grad()
            x = x.cuda()
            t_float = t_float.cuda()
            if xt_normalizer_train is not None:
                x = xt_normalizer_train.normalize(x,t_float, update=False)
            
            t_float_pred = model(x)
            loss = loss_fn(t_float_pred, t_float.type(torch.float32))
            loss.backward()
            loss_arr.append(loss.item())
            bar.set_description(f'Ep:{epoch} loss {np.mean(loss_arr)} val_loss {best_val_loss}')
            optimizer.step()
            if wandb_logger is not None:
                wandb_logger.log_metrics({'train_loss_step':loss.item()})
            
        epoch_losses.append(np.mean(loss_arr))
        scheduler.step(epoch_losses[-1])
        # validation
        model.eval()
        val_losses = []
        for i, (x, t_float) in enumerate(val_loader):
            x = x.cuda()
            t_float = t_float.cuda()
            if xt_normalizer_val is not None:
                x = xt_normalizer_val.normalize(x,t_float, update=False)
            t_float_pred = model(x)
            loss = loss_fn(t_float_pred, t_float.type(torch.float32))
            val_losses.append(loss.item())


        if wandb_logger is not None:
            wandb_logger.log_metrics({'val_loss':np.mean(val_losses)})
        # print(f'Ep:{epoch} Val loss {np.mean(val_losses)}')
        # save best model

        if np.mean(val_losses) < best_val_loss:
            best_val_loss = np.mean(val_losses)
            model_fpath = os.path.join(opt['path']['experiment_root'],'best_time_predictor.pth')
            torch.save(model.state_dict(), model_fpath)
            print('Saved best model', model_fpath)

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--config', type=str, default='config/splitting_hagen_time_predictor.json')
    parser.add_argument('--rootdir', type=str, default='/group/jug/ashesh/training/diffsplit')
    parser.add_argument('-enable_wandb', action='store_true')
    args = parser.parse_args()
    opt = load_json(args.config)
    opt['enable_wandb'] = args.enable_wandb
    experiment_root, expname = get_workdir(opt, args.rootdir, use_max_version=False)

    opt['path']['experiment_root'] = experiment_root
    opt['experiment_name'] = expname
    
    for key, path in opt['path'].items():
        if 'resume' not in key and 'experiments' not in key:
            opt['path'][key] = os.path.join(experiment_root, path)
            mkdirs(opt['path'][key])

    start_training(opt)