train_single_ddp.py

import argparse
import logging
import os
import shutil
import time
import timeit

import numpy as np
import cv2 as cv
cv.setNumThreads(0)
import torch
import torch.backends.cudnn as cudnn
import torch.distributed as torch_dist
import torch.nn.functional as F
from torch import nn, optim
from torch.utils import data
from torchvision.utils import save_image
from tqdm import tqdm

from config import get_cfg_defaults
from dataset.VMD import VideoMattingDataset
from models.model import FullModel
from utils.utils import OPT_DICT, STR_DICT, \
    AverageMeter, create_logger, torch_barrier, reduce_tensor

def write_image(outdir, out, step, max_batch=4):
    with torch.no_grad():
        scaled_imgs, scaled_tris, alphas, comps, gts, fgs, bgs = out
        b, s, _, h, w = scaled_imgs.shape
        b = max_batch if b > max_batch else b
        save_image(scaled_imgs[:max_batch].reshape(b*s, 3, h, w), os.path.join(outdir, 'vis_image_{}.png'.format(step)), nrow=s)
        save_image(scaled_tris[:max_batch].reshape(b*s, 1, h, w), os.path.join(outdir, 'vis_tris_{}.png'.format(step)), nrow=s)
        save_image(alphas[:max_batch].reshape(b*s, 1, h, w), os.path.join(outdir, 'vis_as_{}.png'.format(step)), nrow=s)
        #save_image(refines[:max_batch].reshape(b*s, 1, h, w), os.path.join(outdir, 'vis_refs_{}.png'.format(step)), nrow=s)
        #save_image(trimasks[:max_batch].reshape(b*s, 1, h, w), os.path.join(outdir, 'vis_masks_{}.png'.format(step)), nrow=s)
        save_image(comps[:max_batch].reshape(b*s, 3, h, w), os.path.join(outdir, 'vis_comps_{}.png'.format(step)), nrow=s)
        save_image(gts[:max_batch].reshape(b*s, 1, h, w), os.path.join(outdir, 'vis_gts_{}.png'.format(step)), nrow=s)
        save_image(fgs[:max_batch].reshape(b*s, 3, h, w), os.path.join(outdir, 'vis_fgs_{}.png'.format(step)), nrow=s)
        save_image(bgs[:max_batch].reshape(b*s, 3, h, w), os.path.join(outdir, 'vis_bgs_{}.png'.format(step)), nrow=s)

def train(epoch, trainloader, steps_per_val, base_lr,
          total_epochs, optimizer, model, 
          adjust_learning_rate, print_freq, 
          image_freq, image_outdir, local_rank, sub_losses):    
    # Training
    model.train()
    
    batch_time = AverageMeter()
    ave_loss = AverageMeter()
    tic = time.time()
    cur_iters = epoch*steps_per_val
    for i_iter, dp in enumerate(trainloader):
        def handle_batch():
            fg, bg, a, _ = dp      # [B, 3, 3 or 1, H, W]
            #print (a.shape)
            out = model(a, fg, bg)
            L_alpha = out[0].mean()
            L_comp = out[1].mean()
            L_grad = out[2].mean()
            vis_alpha = L_alpha.detach().item()
            vis_comp = L_comp.detach().item()
            vis_grad = L_grad.detach().item()
            #L_temp = out[3].mean()
            #loss['L_total'] = 0.5 * loss['L_alpha'] + 0.5 * loss['L_comp'] + loss['L_grad'] + 0.5 * loss['L_temp']
            #loss['L_total'] = loss['L_alpha'] + loss['L_comp'] + loss['L_grad'] + loss['L_temp']
            loss = L_alpha + L_comp + L_grad

            model.zero_grad()
            loss.backward()
            optimizer.step()
            return loss.detach(), vis_alpha, vis_comp, vis_grad, out[3:]

        loss, vis_alpha, vis_comp, vis_grad, vis_images = handle_batch()

        reduced_loss = reduce_tensor(loss)
        # measure elapsed time
        batch_time.update(time.time() - tic)
        tic = time.time()

        # update average loss
        ave_loss.update(reduced_loss.item())
        torch_barrier()

        adjust_learning_rate(optimizer,
                            base_lr,
                            total_epochs * steps_per_val,
                            i_iter+cur_iters)

        if i_iter % print_freq == 0 and local_rank <= 0:
            msg = 'Iter:[{}/{}], Time: {:.2f}, '.format(\
                i_iter+cur_iters, total_epochs * steps_per_val, batch_time.average())
            msg += 'lr: {}, Avg. Loss: {:.6f} | Current: Loss: {:.6f}, '.format(
                [x['lr'] for x in optimizer.param_groups],
                ave_loss.average(), ave_loss.value())
            msg += '{}: {:.4f} {}: {:.4f} {}: {:.4f}'.format(
                sub_losses[0], vis_alpha, 
                sub_losses[1], vis_comp,
                sub_losses[2], vis_grad)
            logging.info(msg)
        
        if i_iter % image_freq == 0 and local_rank <= 0:
            write_image(image_outdir, vis_images, i_iter+cur_iters)
            
def validate(testloader, model, test_size, local_rank):
    if local_rank <= 0:
        logging.info('Start evaluation...')
    model.eval()
    ave_loss = AverageMeter()
    with torch.no_grad():
        iterator = tqdm(testloader, ascii=True) if local_rank <= 0 else testloader
        for batch in iterator:
            def handle_batch():
                fg, bg, a, _ = batch      # [B, 3, 3 or 1, H, W]
                out = model(a, fg, bg)
                L_alpha = out[0].mean()
                L_comp = out[1].mean()
                L_grad = out[2].mean()
                #L_temp = out[3].mean()
                #loss['L_total'] = 0.5 * loss['L_alpha'] + 0.5 * loss['L_comp'] + loss['L_grad'] + 0.5 * loss['L_temp']
                #loss['L_total'] = loss['L_alpha'] + loss['L_comp'] + loss['L_grad'] + loss['L_temp']
                loss = L_alpha + L_comp + L_grad
                return loss

            loss = handle_batch()
            reduced_loss = reduce_tensor(loss)

            ave_loss.update(reduced_loss.detach().item())
    if local_rank <= 0:
        logging.info('Validation loss: {:.6f}'.format(ave_loss.average()))
    return ave_loss.average()
    #logging.info('Validation loss: {:.6f}, E_loss: {:.6f}, O_loss: {:.6f} A_loss: {:.6f}'.format(
    #              ave_loss.average(), ave_eloss.average(), ave_oloss.average(), ave_aloss.average()))
    #return ave_loss

def get_sampler(dataset, shuffle=True):
    if torch_dist.is_initialized():
        from torch.utils.data.distributed import DistributedSampler
        return DistributedSampler(dataset, shuffle=shuffle)
    else:
        return None

def main(cfg_name, cfg, local_rank):
    #torch.autograd.set_detect_anomaly(True)
    cfg_name += cfg.SYSTEM.EXP_SUFFIX
    random_seed = cfg.SYSTEM.RANDOM_SEED
    #assert local_rank >= 0
    load_ckpt = cfg.TRAIN.LOAD_CKPT
    base_lr = cfg.TRAIN.BASE_LR
    weight_decay = cfg.TRAIN.WEIGHT_DECAY
    output_dir = cfg.SYSTEM.OUTDIR
    start = timeit.default_timer()
    # cudnn related setting
    cudnn.benchmark = cfg.SYSTEM.CUDNN_BENCHMARK
    cudnn.deterministic = cfg.SYSTEM.CUDNN_DETERMINISTIC
    cudnn.enabled = cfg.SYSTEM.CUDNN_ENABLED
    if random_seed > 0:
        import random
        if local_rank <= 0:
            print('Seeding with', random_seed)
        random.seed(random_seed+local_rank)
        torch.manual_seed(random_seed+local_rank)    

    if local_rank >= 0:
        device = torch.device('cuda:{}'.format(local_rank))    
        torch.cuda.set_device(device)
        torch.distributed.init_process_group(
            backend="nccl", init_method="env://",
        )    
    else:
        device = torch.device('cuda:0')
        torch.cuda.set_device(device)

    if local_rank <= 0:
        logger, final_output_dir = create_logger(output_dir, cfg_name, 'train')
        print (cfg)
        with open(os.path.join(final_output_dir, 'config.yaml'), 'w') as f:
            f.write(str(cfg))
        image_outdir = os.path.join(final_output_dir, 'training_images')
        os.makedirs(os.path.join(final_output_dir, 'training_images'), exist_ok=True)
    else:
        image_outdir = None

    # build model
    model = FullModel(model=cfg.MODEL, agg_window=cfg.AGG_WINDOW, eps=0, \
        freeze_backbone=cfg.TRAIN.FREEZE_BACKBONE)
    if local_rank <= 0 and cfg.MODEL[4:7] == 'res':
        encoder_weight_name = 'resnet34-333f7ec4.pth' \
            if cfg.MODEL[7:9] == '34' else 'resnext50_32x4d-7cdf4587.pth'
        if os.path.exists(encoder_weight_name):
            encoder_weight = torch.load(encoder_weight_name, map_location='cpu')
            dummy = torch.zeros_like(encoder_weight['conv1.weight'][:, :1, :, :])
            encoder_weight['conv1.weight'] = torch.cat([encoder_weight['conv1.weight'], dummy], dim=1)
            missing_keys, unexpected_keys = model.NET.encoder.load_state_dict(encoder_weight, strict=False)
            logger.info('Missing keys: ' + str(sorted(missing_keys)))
            logger.info('Unexpected keys: ' + str(sorted(unexpected_keys)))
            logger.info("=> loaded checkpoint from {} ImageNet pretrain weight.".format(encoder_weight_name))
    torch_barrier()

    # prepare data
    train_dataset = VideoMattingDataset(
        data_root=cfg.DATASET.PATH,
        image_shape=cfg.TRAIN.TRAIN_INPUT_SIZE,
        mode='train',
        use_subset=cfg.DATASET.SUBSET,
        plus1=cfg.MODEL.startswith('vmn_res'),
        no_flow=True,
        sample_length=3
    )
    train_sampler = get_sampler(train_dataset)
    if local_rank < 0:
        trainloader = torch.utils.data.DataLoader(
            train_dataset,
            batch_size=cfg.TRAIN.BATCH_SIZE_PER_GPU,
            shuffle=True,
            num_workers=cfg.SYSTEM.NUM_WORKERS,
            pin_memory=True,
            drop_last=True)
    else:
        trainloader = torch.utils.data.DataLoader(
            train_dataset,
            batch_size=cfg.TRAIN.BATCH_SIZE_PER_GPU,
            #shuffle=True,
            num_workers=cfg.SYSTEM.NUM_WORKERS,
            pin_memory=True,
            drop_last=True,
            sampler=train_sampler)
    test_dataset = VideoMattingDataset(
        data_root=cfg.DATASET.PATH,
        image_shape=cfg.TRAIN.VAL_INPUT_SIZE,
        mode='val',
        use_subset=cfg.DATASET.SUBSET,
        plus1=cfg.MODEL.startswith('vmn_res'),
        no_flow=True,           # to accelerate validation during training we don't compute temporal loss
        sample_length=3,        # we also reduce the sample_length to 3
    )
    testloader = torch.utils.data.DataLoader(
        test_dataset,
        batch_size=cfg.TRAIN.VAL_BATCH_SIZE_PER_GPU,
        shuffle=False,
        num_workers=cfg.SYSTEM.NUM_WORKERS,
        pin_memory=True,
        drop_last=False,
        sampler=get_sampler(test_dataset, shuffle=False)
    )

    if load_ckpt != '':
        dct = torch.load(load_ckpt, map_location=torch.device('cpu'))
        # model_dict = model.state_dict()
        # pretrained_dict = {k: v for k, v in dct.items()
        #             if k in model_dict.keys()}
        # model.load_state_dict(pretrained_dict)
        missing_keys, unexpected_keys = model.NET.load_state_dict(dct, strict=False)
        if local_rank <= 0:
            logger.info('Missing keys: ' + str(sorted(missing_keys)))
            logger.info('Unexpected keys: ' + str(sorted(unexpected_keys)))
            logger.info("=> loaded checkpoint from {}".format(load_ckpt))
        torch_barrier()

    if local_rank >= 0:
        # FBA particularly uses batch_size == 1, thus no syncbn here
        if not cfg.MODEL.endswith('fba'):
            model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model)
        model = model.to(device)
        model = torch.nn.parallel.DistributedDataParallel(
            model,
            find_unused_parameters=True,
            device_ids=[local_rank],
            output_device=local_rank
        )
    else:
        model = torch.nn.DataParallel(model, device_ids=[device])

    # optimizer
    if cfg.TRAIN.FREEZE_BACKBONE:
        params_dict = {k: v for k, v in model.named_parameters() \
            if v.requires_grad and (k[11:] in missing_keys and not k[11:].startswith('semantic'))}
    else:
        params_dict = {k: v for k, v in model.named_parameters() \
            if v.requires_grad and not k[11:].startswith('semantic')}
        
    params_count = 0
    if local_rank <= 0:
        logging.info('=> Parameters needs to be optimized:')
        for k in sorted(params_dict):
            logging.info('\t=> {}, size: {}'.format(k, list(params_dict[k].size())))
            params_count += params_dict[k].shape.numel()
        logging.info('=> Total Parameters: {}'.format(params_count))
        
    params = [{'params': list(params_dict.values()), 'lr': base_lr}]
    optimizer = OPT_DICT[cfg.TRAIN.OPTIMIZER](params, lr=base_lr, weight_decay=weight_decay)
    adjust_lr = STR_DICT[cfg.TRAIN.LR_STRATEGY]

    total_steps = cfg.TRAIN.TOTAL_STEPS
    steps_per_val = len(trainloader)
    print_freq = cfg.TRAIN.PRINT_FREQ
    image_freq = cfg.TRAIN.IMAGE_FREQ
    #assert total_steps % steps_per_val == 0
    #assert steps_per_val % print_freq == 0
    
    #validate(testloader, model, len(test_dataset), local_rank)
    sub_losses = ['L_alpha', 'L_comp', 'L_grad'] if not cfg.MODEL.endswith('fba') else \
                 ['L_alpha_comp', 'L_lap', 'L_grad']
    best_loss = 1e+8
    for epoch in range(total_steps):
        if train_sampler is not None:
            train_sampler.set_epoch(epoch)
        train(epoch, trainloader, steps_per_val, base_lr, total_steps,
              optimizer, model, adjust_lr, print_freq, image_freq, \
              image_outdir, local_rank, sub_losses)
        torch_barrier()
        if epoch >= 15:
            val_loss = validate(testloader, model, len(test_dataset), local_rank)
        else:
            val_loss = best_loss
        torch_barrier()

        if local_rank <= 0:
            weight_fn = os.path.join(final_output_dir,\
                'checkpoint_{}.pth.tar'.format(epoch+1))
            logger.info('=> saving checkpoint to {}'.format(weight_fn))
            torch.save(model.module.NET.state_dict(), weight_fn)
            if val_loss < best_loss:
                best_loss = val_loss
                shutil.copyfile(weight_fn, os.path.join(final_output_dir, 'best.pth'))
                logger.info('=> new minimum loss. copy to best.pth')
        
    end = timeit.default_timer()
    torch_barrier()
    if local_rank <= 0:
        logger.info('Time: %d sec.' % np.int((end-start)))
        logger.info('Done')


def parse_args():
    parser = argparse.ArgumentParser(description='Train network')
    
    parser.add_argument('--cfg',
                        help='experiment configure file name',
                        required=True,
                        type=str)
    parser.add_argument("--local_rank", type=int, default=-1)
    parser.add_argument('opts',
                        help="Modify config options using the command-line",
                        default=None,
                        nargs=argparse.REMAINDER)

    args = parser.parse_args()

    cfg = get_cfg_defaults()
    cfg.merge_from_file(args.cfg)
    cfg.merge_from_list(args.opts)
    cfg.freeze()

    return args, cfg

if __name__ == "__main__":
    args, cfg = parse_args()
    main(os.path.splitext(os.path.basename(args.cfg))[0], cfg, args.local_rank)