train_rfda_ft.py

import os
os.environ['CUDA_VISIBLE_DEVICES'] = "0,1"
import math
import time
import yaml
import argparse
import random
import torch
import torch.optim as optim
import torch.nn as nn
import os.path as op
import numpy as np
from tqdm import tqdm
from torch.nn.parallel import DistributedDataParallel as DDP
from collections import OrderedDict
from PIL import Image
import utils  # my tool box
import dataset
from net_rfda import RFDA
# torch.autograd.set_detect_anomaly(True)
def receive_arg():
    """Process all hyper-parameters and experiment settings.
    
    Record in opts_dict."""
    parser = argparse.ArgumentParser()
    parser.add_argument(
        '--opt_path', type=str, default='option_R3_mfqev2_2G.yml', 
        help='Path to option YAML file.'
        )
    parser.add_argument(
        '--local_rank', type=int, default=0, 
        help='Distributed launcher requires.'
        )
    args = parser.parse_args()
    
    with open(args.opt_path, 'r') as fp:
        opts_dict = yaml.load(fp, Loader=yaml.FullLoader)

    opts_dict['opt_path'] = args.opt_path
    opts_dict['train']['rank'] = args.local_rank

    if opts_dict['train']['exp_name'] == None:
        opts_dict['train']['exp_name'] = utils.get_timestr()

    opts_dict['train']['log_path'] = op.join(
        "exp", opts_dict['train']['exp_name'], "log.log"
        )
    opts_dict['train']['checkpoint_save_path_pre'] = op.join(
        "exp", opts_dict['train']['exp_name'], "ckp_"
        )
    
    opts_dict['train']['num_gpu'] = torch.cuda.device_count()
    if opts_dict['train']['num_gpu'] > 1:
        opts_dict['train']['is_dist'] = True
    else:
        opts_dict['train']['is_dist'] = False
    
    opts_dict['test']['restore_iter'] = int(
        opts_dict['test']['restore_iter']
        )

    return opts_dict

def get_lr(lr,milestones,it,gamma):
    count=0
    for milestone in milestones:
        if(it>milestone):count+=1
    return lr*pow(gamma,count)

def main():
    # ==========
    # parameters
    # ==========

    opts_dict = receive_arg()
    rank = opts_dict['train']['rank']
    unit = opts_dict['train']['criterion']['unit']
    num_iter = int(opts_dict['train']['num_iter'])
    interval_print = int(opts_dict['train']['interval_print'])
    interval_val = int(opts_dict['train']['interval_val'])
    
    # ==========
    # init distributed training
    # ==========
    # opts_dict['train']['is_dist'] = False
    if opts_dict['train']['is_dist']:
        utils.init_dist(
            local_rank=rank, 
            backend='nccl'
            )

    # TO-DO: load resume states if exists
    pass

    # ==========
    # create logger
    # ==========

    if rank == 0:
        log_dir = op.join("exp", opts_dict['train']['exp_name'])
        if not os.path.exists(log_dir):
            utils.mkdir(log_dir)
        log_fp = open(opts_dict['train']['log_path'], 'a')

        # log all parameters
        msg = (
            f"{'<' * 10} Hello {'>' * 10}\n"
            f"Timestamp: [{utils.get_timestr()}]\n"
            f"\n{'<' * 10} Options {'>' * 10}\n"
            f"{utils.dict2str(opts_dict)}"
            )
        print(msg)
        log_fp.write(msg + '\n')
        log_fp.flush()

    # ==========
    # TO-DO: init tensorboard
    # ==========

    pass
    
    # ==========
    # fix random seed
    # ==========

    seed = opts_dict['train']['random_seed']
    # >I don't know why should rs + rank
    utils.set_random_seed(seed + rank)

    # ========== 
    # Ensure reproducibility or Speed up
    # ==========

    #torch.backends.cudnn.benchmark = False  # if reproduce
    #torch.backends.cudnn.deterministic = True  # if reproduce
    torch.backends.cudnn.benchmark = True  # speed up

    # ==========
    # create train and val data prefetchers
    # ==========
    
    # create datasets
    train_ds_type = opts_dict['dataset']['train']['type']
    val_ds_type = opts_dict['dataset']['val']['type']
    radius = opts_dict['network']['radius']
    radius_real = 3
    assert train_ds_type in dataset.__all__, \
        "Not implemented!"
    assert val_ds_type in dataset.__all__, \
        "Not implemented!"
    train_ds_cls = getattr(dataset, train_ds_type)
    val_ds_cls = getattr(dataset, val_ds_type)
    train_ds = train_ds_cls(
        opts_dict=opts_dict['dataset']['train'], 
        radius=radius
        )
    val_ds = val_ds_cls(
        opts_dict=opts_dict['dataset']['val'], 
        radius=radius
        )

    # create datasamplers
    train_sampler = utils.DistSampler(
        dataset=train_ds, 
        num_replicas=opts_dict['train']['num_gpu'], 
        rank=rank, 
        ratio=opts_dict['dataset']['train']['enlarge_ratio']
        )
    val_sampler = None  # no need to sample val data

    # create dataloaders
    train_loader = utils.create_dataloader(
        dataset=train_ds, 
        opts_dict=opts_dict, 
        sampler=train_sampler, 
        phase='train',
        seed=opts_dict['train']['random_seed']
        )
    val_loader = utils.create_dataloader(
        dataset=val_ds, 
        opts_dict=opts_dict, 
        sampler=val_sampler, 
        phase='val'
        )
    assert train_loader is not None

    batch_size = opts_dict['dataset']['train']['batch_size_per_gpu'] * \
        opts_dict['train']['num_gpu']  # divided by all GPUs
    num_iter_per_epoch = math.ceil(len(train_ds) * \
        opts_dict['dataset']['train']['enlarge_ratio'] / batch_size)
    num_epoch = math.ceil(num_iter / num_iter_per_epoch)
    val_num = len(val_ds)
    
    # create dataloader prefetchers
    tra_prefetcher = utils.CPUPrefetcher(train_loader)
    val_prefetcher = utils.CPUPrefetcher(val_loader)

    # ==========
    # create model
    # ==========

    model = RFDA(opts_dict=opts_dict['network'])

    model = model.to(rank)
    # if opts_dict['train']['is_dist']:
    #     model = DDP(model, device_ids=[rank])

    # ==========
    # define loss func & optimizer & scheduler & criterion
    # ==========

    # define loss func
    if opts_dict['train']['loss']['type'] == 'CharbonnierLoss':
        print("CharbonnierLoss")
        loss_func = utils.CharbonnierLoss(**opts_dict['train']['loss'])
    elif opts_dict['train']['loss']['type'] == 'L1':
        print("L1Loss")
        loss_func = nn.L1Loss()
    elif opts_dict['train']['loss']['type'] == 'L2':
        print("MSELoss")
        loss_func = nn.MSELoss()

    # TODO 
    # loss_func = torch.nn.MSELoss()
    # define optimizer
    assert opts_dict['train']['optim'].pop('type') == 'Adam', \
        "Not implemented."
    # print(opts_dict['train']['optim'],"???")

    # optimizer = optim.Adam(model.parameters(),**opts_dict['train']['optim'])
    base_params = filter(lambda p: id(p) not in list(map(id, model.qenet.parameters()))+list(map(id, model.ffnet.parameters())),model.parameters())
    optimizer = optim.Adam(
        [
            {"params": base_params},
            {"params": model.qenet.parameters(), "lr": opts_dict['train']['optim']["lr"] / 10.0},
            {"params": model.ffnet.parameters(), "lr": opts_dict['train']['optim']["lr"] / 10.0},
        ],
        **opts_dict['train']['optim']
        )
    # define scheduler
    milestones=[]
    for milestone in opts_dict['train']['scheduler']['milestones']:
        milestones.append(int(milestone*num_iter))
    gamma=opts_dict['train']['scheduler']['gamma']
    opt_lr=opts_dict['train']['optim']['lr']
    set_milestones=set(milestones)

    # define criterion
    assert opts_dict['train']['criterion'].pop('type') == \
        'PSNR', "Not implemented."
    criterion = utils.PSNR()
    start_iter = 0
    # ==========
    # TO-DO: resume training & load pre-trained model
    # ==========

    
    filename = '/remote-home/myzhao/MM_CKPS/R3_QP37.pt'
    print("=> loading checkpoint '{}'".format(filename))
    checkpoint = torch.load(filename, map_location=torch.device('cpu'))
    
    # start_iter = checkpoint['num_iter_accum']
    # optimizer.load_state_dict(checkpoint['optimizer'])
    # if opts_dict['train']['scheduler']['is_on']:
    #     for i in range(start_iter):
    #         scheduler.step()
    # 第一阶段先load一下原有模型
    if 'module.' in list(checkpoint['state_dict'].keys())[0]:  # multi-gpu training
        new_state_dict = OrderedDict()
        for k, v in checkpoint['state_dict'].items():
            name = k[7:]  # remove module
            new_state_dict[name] = v
        model.load_state_dict(new_state_dict,strict=False)
    else:  # single-gpu training
        model.load_state_dict(checkpoint['state_dict'])
    model = DDP(model, device_ids=[rank],find_unused_parameters=True)

    now_lr=get_lr(opt_lr,milestones,start_iter,gamma)
    # print(opt_lr,milestones,start_iter,gamma)
    # print("now_lr",now_lr)
    # if (opts_dict["network"]['qenet']['netname'] != 'C2CNET'):
    #     for param_group in optimizer.param_groups:
    #         param_group["lr"] = now_lr
    optimizer.param_groups[0]["lr"] = now_lr
    optimizer.param_groups[1]["lr"] = now_lr/10.0
    optimizer.param_groups[2]["lr"] = now_lr/10.0
    print("=> loaded checkpoint Success ")
    
    torch.autograd.set_detect_anomaly(True)
    # start_iter = 0  # should be restored
    start_epoch = start_iter // num_iter_per_epoch

    # display and log
    if rank == 0:
        msg = (
            f"\n{'<' * 10} Dataloader {'>' * 10}\n"
            f"total iters: [{num_iter}]\n"
            f"total epochs: [{num_epoch}]\n"
            f"iter per epoch: [{num_iter_per_epoch}]\n"
            f"val sequence: [{val_num}]\n"
            f"start from iter: [{start_iter}]\n"
            f"start from epoch: [{start_epoch}]"
            )
        print(msg)
        log_fp.write(msg + '\n')
        log_fp.flush()

    # Test LR
    # for i in range(num_iter):
    #     optimizer.step()
    #     scheduler.step()
    #     lr = optimizer.param_groups[0]['lr']
    #     if i%1000==0:
    #         print(i,lr,i/num_iter)
    # os._exit(233)
    # ==========
    # evaluate original performance, e.g., PSNR before enhancement
    # ==========
    vid_num = val_ds.get_vid_num()
    if opts_dict['train']['pre-val'] and rank == 0:
        msg = f"\n{'<' * 10} Pre-evaluation {'>' * 10}"
        print(msg)
        log_fp.write(msg + '\n')

        per_aver_dict = {}
        for i in range(vid_num):
            per_aver_dict[i] = utils.Counter()
        pbar = tqdm(
                total=val_num, 
                ncols=opts_dict['train']['pbar_len']
                )

        # fetch the first batch
        val_prefetcher.reset()
        val_data = val_prefetcher.next()

        while val_data is not None:
            # get data
            gt_data = val_data['gt'].to(rank)  # (B [RGB] H W)
            lq_data = val_data['lq'].to(rank)  # (B T [RGB] H W)
            index_vid = val_data['index_vid'].item()
            name_vid = val_data['name_vid'][0]  # bs must be 1!
            b, _, _, _, _  = lq_data.shape
            # 只算Y通道的！！！
            # lq_data = lq_data[:,:,0,...]*0.299 + lq_data[:,:,1,...]*0.587 + lq_data[:,:,2,...]*0.114
            # gt_data = gt_data[:,0,...]*0.299 + gt_data[:,1,...]*0.587 + gt_data[:,2,...]*0.114
            # eval
            batch_perf = np.mean(
                [criterion(lq_data[i,radius,...], gt_data[i]) for i in range(b)]
                )  # bs must be 1!
            
            # log
            per_aver_dict[index_vid].accum(volume=batch_perf)

            # display
            pbar.set_description(
                "{:s}: [{:.3f}] {:s}".format(name_vid, batch_perf, unit)
                )
            pbar.update()

            # fetch next batch
            val_data = val_prefetcher.next()

        pbar.close()

        # log
        ave_performance = np.mean([
            per_aver_dict[index_vid].get_ave() for index_vid in range(vid_num)
            ])
        msg = "> ori performance: [{:.3f}] {:s}".format(ave_performance, unit)
        print(msg)
        log_fp.write(msg + '\n')
        log_fp.flush()

    if opts_dict['train']['is_dist']:
        torch.distributed.barrier()  # all processes wait for ending

    if rank == 0:
        msg = f"\n{'<' * 10} Training {'>' * 10}"
        print(msg)
        log_fp.write(msg + '\n')

        # create timer
        total_timer = utils.Timer()  # total tra + val time of each epoch

    # ==========
    # start training + validation (test)
    # ==========

    model.train()
    num_iter_accum = start_iter
    for current_epoch in range(start_epoch, num_epoch + 1):
        # shuffle distributed subsamplers before each epoch
        if opts_dict['train']['is_dist']:
            train_sampler.set_epoch(current_epoch)

        # fetch the first batch
        tra_prefetcher.reset()
        train_data = tra_prefetcher.next()

        # train this epoch
        while train_data is not None:

            # over sign
            num_iter_accum += 1
            if num_iter_accum > num_iter:
                break

            # get data
            gt_data = train_data['gt'].to(rank)  # (B 3 [RGB] H W)
            lq_data = train_data['lq'].to(rank)  # (B T [RGB] H W)
            # print(gt_data.size(),"vs",lq_data.size())
            b, t, c, _, _  = lq_data.shape
            # 重新装载
            if not opts_dict['network']['stdf'].__contains__('netname') or opts_dict['network']['stdf']['netname']=='default':
                input_data = torch.cat(
                    [lq_data[:,:,i,...] for i in range(c)], 
                    dim=1
                    )  # B [R1 ... R7 G1 ... G7 B1 ... B7] H W
            else:
                input_data = lq_data
            loss_sum = 0
            enhanced_datas = []
            for i in range(t):
                # print("i=",i,'t=',t)
                neighbor_list = list(range(i - radius_real, i + radius_real + 1))
                neighbor_list = list(np.clip(neighbor_list, 0, t - 1))
                frm_list = []
                # print(neighbor_list,'vs',lq_data.size())
                for _i in neighbor_list:
                    # print("i=",i)
                    # print(lq_data[:,int(i),:...].size())
                    frm_list.append(lq_data[:,int(_i),...])
                frm_list = torch.cat(frm_list,1).to(rank)
                if i==0:
                    enhanced_data,hint = model(frm_list)
                    enhanced_datas.append(enhanced_data)
                    # loss = torch.mean(torch.stack(
                    #     [loss_func(enhanced_data[j], gt_data[j,i,...]) for j in range(b)]
                    #     ))
                else:
                    enhanced_data,hint = model(frm_list,hint)
                    enhanced_datas.append(enhanced_data)
                    # loss = loss + torch.mean(torch.stack(
                    #     [loss_func(enhanced_data[j], gt_data[j,i,...]) for j in range(b)]
                    #     ))
            optimizer.zero_grad()  # zero grad
            # 计算损失
            loss = 0
            for i in range(t):
                loss = loss + torch.mean(torch.stack(
                    [loss_func(enhanced_datas[i][j,...], gt_data[j,i,...]) for j in range(b)]
                    ))
            loss.backward()  # cal grad
            optimizer.step()  # update parameters
            # os._exit(233)
            # update learning rate
            
            if (num_iter_accum % interval_print == 0) and (rank == 0):
                now_lr = get_lr(opt_lr, milestones, num_iter_accum, gamma)
                if (now_lr != optimizer.param_groups[0]['lr']):
                    # for param_group in optimizer.param_groups:
                    #     param_group["lr"] = now_lr
                    optimizer.param_groups[0]["lr"] = now_lr
                    optimizer.param_groups[1]["lr"] = now_lr/10.0
                    optimizer.param_groups[2]["lr"] = now_lr/10.0
                # display & log
                lr = optimizer.param_groups[0]['lr']
                loss_item = loss.item()
                msg = (
                    f"iter: [{num_iter_accum}]/{num_iter}, "
                    f"epoch: [{current_epoch}]/{num_epoch - 1}, "
                    "lr: [{:.3f}]x1e-4, loss: [{:.4f}]".format(
                        lr*1e4, loss_item
                        )
                    )
                print(msg)
                log_fp.write(msg + '\n')
            
            if ((num_iter_accum % interval_val == 0) or \
                (num_iter_accum == num_iter)) and (rank == 0):
                # save model
                checkpoint_save_path = (
                    f"{opts_dict['train']['checkpoint_save_path_pre']}"
                    f"{num_iter_accum}"
                    ".pt"
                    )
                state = {
                    'num_iter_accum': num_iter_accum, 
                    'state_dict': model.state_dict(),
                    'optimizer': optimizer.state_dict(), 
                    }
                torch.save(state, checkpoint_save_path)
                # print("continue")
                # continuecontinue
                # validation
                with torch.no_grad():
                    per_aver_dict = {}
                    for index_vid in range(vid_num):
                        per_aver_dict[index_vid] = utils.Counter()
                    pbar = tqdm(
                            total=7980, 
                            ncols=opts_dict['train']['pbar_len']
                            )
                
                    # train -> eval
                    model.eval()

                    # fetch the first batch
                    val_prefetcher.reset()
                    val_data = val_prefetcher.next()
                    
                    while val_data is not None:
                        # get data
                        gt_data = val_data['gt']  # (B T [RGB] H W)
                        lq_data = val_data['lq']  # (B T [RGB] H W)
                        # print(gt_data.size(),'vs',lq_data.size())
                        # print(gt_data)
                        # os._exit(2333)
                        index_vid = val_data['index_vid'].item()
                        name_vid = val_data['name_vid'][0]  # bs must be 1!
                        b, t, c, _, _  = lq_data.shape
                        # input_data = torch.cat(
                        #     [lq_data[:,:,i,...] for i in range(c)], 
                        #     dim=1
                        #     )  # B [R1 ... R7 G1 ... G7 B1 ... B7] H W
                        for i in range(t):
                            neighbor_list = list(range(i - radius_real, i + radius_real + 1))
                            neighbor_list = list(np.clip(neighbor_list, 0, t - 1))
                            frm_list = []
                            # print(neighbor_list,'vs',lq_data.size())
                            for _i in neighbor_list:
                                # print("i=",i)
                                # print(lq_data[:,int(i),:...].size())
                                frm_list.append(lq_data[:,int(_i),...])
                            frm_list = torch.cat(frm_list,1).to(rank)
                            if i==0:
                                # print("first size",first.size())
                                enhanced_data,hint = model(frm_list)  # (B [RGB] H W)
                            else:
                                enhanced_data,hint = model(frm_list,hint)
                            # 注意了，接下来我们只算Y通道的，为了公平的比较.
                            # enhanced_data = enhanced_data[:,0,...]*0.299 + enhanced_data[:,1,...]*0.587 + enhanced_data[:,2,...]*0.114
                            # gt_data = gt_data[:,0,...]*0.299 + gt_data[:,1,...]*0.587 + gt_data[:,2,...]*0.114
                            # eval
                            batch_perf = np.mean(
                                [criterion(enhanced_data[j], gt_data[j,i,...].to(rank)) for j in range(b)]
                                ) # bs must be 1!

                            # display
                            pbar.set_description(
                                "{:s}: [{:.3f}] {:s}"
                                .format(name_vid, batch_perf, unit)
                                )
                            pbar.update()

                            # log
                            per_aver_dict[index_vid].accum(volume=batch_perf)
                            # torch.cuda.empty_cache()
                        # fetch next batch
                        torch.cuda.empty_cache()
                        val_data = val_prefetcher.next()
                        # return
                    # end of val
                    pbar.close()

                    # eval -> train
                    model.train()

                # log
                ave_per = np.mean([
                    per_aver_dict[index_vid].get_ave() for index_vid in range(vid_num)
                    ])
                msg = (
                    "> model saved at {:s}\n"
                    "> ave val per: [{:.3f}] {:s}"
                    ).format(
                        checkpoint_save_path, ave_per, unit
                        )
                print(msg)
                log_fp.write(msg + '\n')
                log_fp.flush()

            if opts_dict['train']['is_dist']:
                torch.distributed.barrier()  # all processes wait for ending

            # fetch next batch
            train_data = tra_prefetcher.next()

        # end of this epoch (training dataloader exhausted)

    # end of all epochs

    # ==========
    # final log & close logger
    # ==========

    if rank == 0:
        total_time = total_timer.get_interval() / 3600
        msg = "TOTAL TIME: [{:.1f}] h".format(total_time)
        print(msg)
        log_fp.write(msg + '\n')
        
        msg = (
            f"\n{'<' * 10} Goodbye {'>' * 10}\n"
            f"Timestamp: [{utils.get_timestr()}]"
            )
        print(msg)
        log_fp.write(msg + '\n')
        
        log_fp.close()


if __name__ == '__main__':
    main()