train.py

import os
import cv2 
import time
import random 
import datetime
import argparse
import numpy as np
from tqdm import tqdm
from piq import ssim,psnr
from itertools import cycle

import torch
import torch.nn as nn
from torch.utils import data
import torch.distributed as dist 
from torch.utils.data.distributed import DistributedSampler 
from torch.nn.parallel import DistributedDataParallel as DDP 


from utils import dict2string,mkdir,get_lr,torch2cvimg,second2hours
from loaders import docres_loader
from models import restormer_arch


def seed_torch(seed=1029):
	random.seed(seed)
	os.environ['PYTHONHASHSEED'] = str(seed) 
	np.random.seed(seed)
	torch.manual_seed(seed)
	torch.cuda.manual_seed(seed)
	torch.cuda.manual_seed_all(seed) 
	torch.backends.cudnn.benchmark = False
	torch.backends.cudnn.deterministic = True
    #torch.use_deterministic_algorithms(True)
# seed_torch()


def getBasecoord(h,w):
    base_coord0 = np.tile(np.arange(h).reshape(h,1),(1,w)).astype(np.float32)
    base_coord1 = np.tile(np.arange(w).reshape(1,w),(h,1)).astype(np.float32)
    base_coord = np.concatenate((np.expand_dims(base_coord1,-1),np.expand_dims(base_coord0,-1)),-1)
    return base_coord

def train(args):

    ## DDP init
    dist.init_process_group(backend='nccl',init_method='env://',timeout=datetime.timedelta(seconds=36000))
    torch.cuda.set_device(args.local_rank)
    device = torch.device('cuda',args.local_rank)
    torch.cuda.manual_seed_all(42)

    ### Log file:
    mkdir(args.logdir)
    mkdir(os.path.join(args.logdir,args.experiment_name))
    log_file_path=os.path.join(args.logdir,args.experiment_name,'log.txt')
    log_file=open(log_file_path,'a')
    log_file.write('\n---------------  '+args.experiment_name+'  ---------------\n')
    log_file.close()

    ### Setup tensorboard for visualization
    if args.tboard:
        writer = SummaryWriter(os.path.join(args.logdir,args.experiment_name,'runs'),args.experiment_name)

    ### Setup Dataloader
    datasets_setting = [
        {'task':'deblurring','ratio':1,'im_path':'/home/jiaxin/Training_Data/DocRes_data/train/deblurring/','json_paths':['/home/jiaxin/Training_Data/DocRes_data/train/deblurring/tdd/train.json']},
        {'task':'dewarping','ratio':1,'im_path':'/home/jiaxin/Training_Data/DocRes_data/train/dewarping/','json_paths':['/home/jiaxin/Training_Data/DocRes_data/train/dewarping/doc3d/train_1_19.json']},
        {'task':'binarization','ratio':1,'im_path':'/home/jiaxin/Training_Data/DocRes_data/train/binarization/','json_paths':['/home/jiaxin/Training_Data/DocRes_data/train/binarization/train.json']},
        {'task':'deshadowing','ratio':1,'im_path':'/home/jiaxin/Training_Data/DocRes_data/train/deshadowing/','json_paths':['/home/jiaxin/Training_Data/DocRes_data/train/deshadowing/train.json']},
        {'task':'appearance','ratio':1,'im_path':'/home/jiaxin/Training_Data/DocRes_data/train/appearance/','json_paths':['/home/jiaxin/Training_Data/DocRes_data/train/appearance/trainv2.json']}
        ]


    ratios = [dataset_setting['ratio'] for dataset_setting in datasets_setting]
    datasets = [docres_loader.DocResTrainDataset(dataset=dataset_setting,img_size=args.im_size) for dataset_setting in datasets_setting]
    trainloaders = [{'task':datasets_setting[i],'loader':data.DataLoader(dataset=datasets[i], sampler=DistributedSampler(datasets[i]), batch_size=args.batch_size, num_workers=2, pin_memory=True,drop_last=True),'iter_loader':iter(data.DataLoader(dataset=datasets[i], sampler=DistributedSampler(datasets[i]), batch_size=args.batch_size, num_workers=2, pin_memory=True,drop_last=True))} for i in range(len(datasets))]


    ### test loader
    # for i in tqdm(range(args.total_iter)):
    #     loader_index = random.choices(list(range(len(trainloaders))),ratios)[0]
    #     in_im,gt_im = next(trainloaders[loader_index]['iter_loader'])


    ### Setup Model
    model = restormer_arch.Restormer( 
        inp_channels=6, 
        out_channels=3, 
        dim = 48,
        num_blocks = [2,3,3,4],     
        num_refinement_blocks = 4,
        heads = [1,2,4,8],
        ffn_expansion_factor = 2.66,
        bias = False,
        LayerNorm_type = 'WithBias',   
        dual_pixel_task = True       
    )
    model=DDP(model.cuda(),device_ids=[args.local_rank],output_device=args.local_rank)

    ### Optimizer
    optimizer= torch.optim.AdamW(model.parameters(),lr=args.l_rate,weight_decay=5e-4)

    ### LR Scheduler 
    sched = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, args.total_iter, eta_min=1e-6, last_epoch=-1)

    ### load checkpoint
    iter_start=0
    if args.resume is not None:                                         
        print("Loading model and optimizer from checkpoint '{}'".format(args.resume))
        x = checkpoint['model_state']
        model.load_state_dict(x,strict=False)
        iter_start=checkpoint['iter']
        print("Loaded checkpoint '{}' (iter {})".format(args.resume, iter_start))

    ###-----------------------------------------Training-----------------------------------------
    ##initialize
    scaler = torch.cuda.amp.GradScaler()
    loss_dict = {}
    total_step = 0
    l2 = nn.MSELoss()
    l1 = nn.L1Loss()
    ce = nn.CrossEntropyLoss()
    bce = nn.BCEWithLogitsLoss()
    m = nn.Sigmoid()
    best = 0
    best_ce = 999

    ## total_steps
    for iters in range(iter_start,args.total_iter):
        start_time = time.time()
        loader_index = random.choices(list(range(len(trainloaders))),ratios)[0]

        try:
            in_im,gt_im = next(trainloaders[loader_index]['iter_loader'])
        except StopIteration:
            trainloaders[loader_index]['iter_loader']=iter(trainloaders[loader_index]['loader'])
            in_im,gt_im = next(trainloaders[loader_index]['iter_loader'])
        in_im = in_im.float().cuda()
        gt_im = gt_im.float().cuda()

        binarization_loss,appearance_loss,dewarping_loss,deblurring_loss,deshadowing_loss = 0,0,0,0,0
        with torch.cuda.amp.autocast():
            pred_im = model(in_im,trainloaders[loader_index]['task']['task'])
            if trainloaders[loader_index]['task']['task'] == 'binarization':
                gt_im = gt_im.long()
                binarization_loss = ce(pred_im[:,:2,:,:], gt_im[:,0,:,:])
                loss = binarization_loss
            elif trainloaders[loader_index]['task']['task'] == 'dewarping':
                dewarping_loss = l1(pred_im[:,:2,:,:], gt_im[:,:2,:,:])
                loss = dewarping_loss
            elif trainloaders[loader_index]['task']['task'] == 'appearance':
                appearance_loss = l1(pred_im, gt_im)
                loss = appearance_loss
            elif trainloaders[loader_index]['task']['task'] == 'deblurring':
                deblurring_loss = l1(pred_im, gt_im)
                loss = deblurring_loss
            elif trainloaders[loader_index]['task']['task'] == 'deshadowing':
                deshadowing_loss = l1(pred_im, gt_im)
                loss = deshadowing_loss

        optimizer.zero_grad()
        scaler.scale(loss).backward()
        scaler.step(optimizer)
        scaler.update() 
    
        loss_dict['dew_loss']=dewarping_loss.item() if isinstance(dewarping_loss,torch.Tensor) else 0
        loss_dict['app_loss']=appearance_loss.item() if isinstance(appearance_loss,torch.Tensor) else 0
        loss_dict['des_loss']=deshadowing_loss.item() if isinstance(deshadowing_loss,torch.Tensor) else 0
        loss_dict['deb_loss']=deblurring_loss.item() if isinstance(deblurring_loss,torch.Tensor) else 0
        loss_dict['bin_loss']=binarization_loss.item() if isinstance(binarization_loss,torch.Tensor) else 0
        end_time = time.time()
        duration = end_time-start_time
        ## log
        if (iters+1) % 10 == 0:
            ## print
            print('iters [{}/{}] -- '.format(iters+1,args.total_iter)+dict2string(loss_dict)+' --lr {:6f}'.format(get_lr(optimizer))+' -- time {}'.format(second2hours(duration*(args.total_iter-iters))))
            ## tbord
            if args.tboard:
                for key,value in loss_dict.items():
                    writer.add_scalar('Train '+key+'/Iterations', value, total_step)
            ## logfile
            with open(log_file_path,'a') as f:
                f.write('iters [{}/{}] -- '.format(iters+1,args.total_iter)+dict2string(loss_dict)+' --lr {:6f}'.format(get_lr(optimizer))+' -- time {}'.format(second2hours(duration*(args.total_iter-iters)))+'\n')


        if (iters+1) % 5000 == 0:
            state = {'iters': iters+1,
                     'model_state': model.state_dict(),
                     'optimizer_state' : optimizer.state_dict(),}
            if not os.path.exists(os.path.join(args.logdir,args.experiment_name)):
                 os.system('mkdir ' + os.path.join(args.logdir,args.experiment_name))
            if torch.distributed.get_rank()==0:
                torch.save(state, os.path.join(args.logdir,args.experiment_name,"{}.pkl".format(iters+1)))

        sched.step()



if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Hyperparams')
    parser.add_argument('--im_size', nargs='?', type=int, default=256, 
                        help='Height of the input image')
    parser.add_argument('--total_iter', nargs='?', type=int, default=100000, 
                        help='# of the epochs')
    parser.add_argument('--batch_size', nargs='?', type=int, default=10, 
                        help='Batch Size')
    parser.add_argument('--l_rate', nargs='?', type=float, default=2e-4, 
                        help='Learning Rate')
    parser.add_argument('--resume', nargs='?', type=str, default=None,    
                        help='Path to previous saved model to restart from')
    parser.add_argument('--logdir', nargs='?', type=str, default='./checkpoints/',    
                        help='Path to store the loss logs')
    parser.add_argument('--tboard', dest='tboard', action='store_true', 
                        help='Enable visualization(s) on tensorboard | False by default')
    parser.add_argument('--local_rank',type=int,default=0,metavar='N')    
    parser.add_argument('--experiment_name', nargs='?', type=str,default='experiment_name',
                        help='the name of this experiment')
    parser.set_defaults(tboard=False)
    args = parser.parse_args()

    train(args)