UnOpticalFlow/train.py

import os, sys
import yaml
sys.path.append(os.path.dirname(os.path.abspath(__file__)))
from core.dataset import KITTI_RAW, KITTI_Prepared, SINTEL_RAW, SINTEL_Prepared, NYU_Prepare, NYU_v2, KITTI_Odo
from core.networks import get_model
from core.config import generate_loss_weights_dict
from core.visualize import Visualizer
from core.evaluation import load_gt_flow_kitti, load_gt_mask
from test import test_kitti_2012, test_kitti_2015, test_eigen_depth, test_nyu, load_nyu_test_data

from collections import OrderedDict
import torch
import torch.utils.data
from tqdm import tqdm
import shutil
import pickle
import pdb
import random
import numpy as np
import torch.backends.cudnn as cudnn


def save_model(iter_, model_dir, filename, model, optimizer):
    torch.save({"iteration": iter_, "model_state_dict": model.state_dict(), 'optimizer_state_dict': optimizer.state_dict()}, os.path.join(model_dir, filename))

def load_model(model_dir, filename, model, optimizer):
    data = torch.load(os.path.join(model_dir, filename))
    iter_ = data['iteration']
    model.load_state_dict(data['model_state_dict'])
    optimizer.load_state_dict(data['optimizer_state_dict'])
    return iter_, model, optimizer

def train(cfg):
    # load model and optimizer
    model = get_model(cfg.mode)(cfg)
    if cfg.multi_gpu:
        model = torch.nn.DataParallel(model)
    model = model.cuda()
    optimizer = torch.optim.Adam([{'params': filter(lambda p: p.requires_grad, model.parameters()), 'lr': cfg.lr}])

    # Load Pretrained Models
    if cfg.resume:
        if cfg.iter_start > 0:
            cfg.iter_start, model, optimizer = load_model(cfg.model_dir, 'iter_{}.pth'.format(cfg.iter_start), model, optimizer)
        else:
            cfg.iter_start, model, optimizer = load_model(cfg.model_dir, 'last.pth', model, optimizer)
    elif cfg.flow_pretrained_model:
        data = torch.load(cfg.flow_pretrained_model)['model_state_dict']
        renamed_dict = OrderedDict()
        for k, v in data.items():
            if cfg.multi_gpu:
                name = 'module.model_flow.' + k
            elif cfg.mode == 'flowposenet':
                name = 'model_flow.' + k
            else:
                name = 'model_pose.model_flow.' + k
            renamed_dict[name] = v
        missing_keys, unexp_keys = model.load_state_dict(renamed_dict, strict=False)
        print(missing_keys)
        print(unexp_keys)
        print('Load Flow Pretrained Model from ' + cfg.flow_pretrained_model)
    if cfg.depth_pretrained_model and not cfg.resume:
        data = torch.load(cfg.depth_pretrained_model)['model_state_dict']
        if cfg.multi_gpu:
            renamed_dict = OrderedDict()
            for k, v in data.items():
                name = 'module.' + k
                renamed_dict[name] = v
            missing_keys, unexp_keys = model.load_state_dict(renamed_dict, strict=False)
        else:
            missing_keys, unexp_keys = model.load_state_dict(data, strict=False)
        print(missing_keys)
        print('##############')
        print(unexp_keys)
        print('Load Depth Pretrained Model from ' + cfg.depth_pretrained_model)
   
    loss_weights_dict = generate_loss_weights_dict(cfg)
    visualizer = Visualizer(loss_weights_dict, cfg.log_dump_dir)

    # load dataset
    data_dir = os.path.join(cfg.prepared_base_dir, cfg.prepared_save_dir)
    if not os.path.exists(os.path.join(data_dir, 'train.txt')):
        if cfg.dataset == 'kitti_depth':
            kitti_raw_dataset = KITTI_RAW(cfg.raw_base_dir, cfg.static_frames_txt, cfg.test_scenes_txt)
            kitti_raw_dataset.prepare_data_mp(data_dir, stride=1)
        elif cfg.dataset == 'sintel_raw':
            sintel_raw_dataset = SINTEL_RAW(cfg.raw_base_dir)
            sintel_raw_dataset.prepare_data_mp(data_dir, cfg.stride)
        elif cfg.dataset == 'kitti_odo':
            kitti_raw_dataset = KITTI_Odo(cfg.raw_base_dir)
            kitti_raw_dataset.prepare_data_mp(data_dir, stride=1)
        elif cfg.dataset == 'nyuv2':
            nyu_raw_dataset = NYU_Prepare(cfg.raw_base_dir, cfg.nyu_test_dir)
            nyu_raw_dataset.prepare_data_mp(data_dir, stride=10)
        else:
            raise NotImplementedError
        
    
    if cfg.dataset == 'kitti_depth':
        dataset = KITTI_Prepared(data_dir, num_scales=cfg.num_scales, img_hw=cfg.img_hw, num_iterations=(cfg.num_iterations - cfg.iter_start) * cfg.batch_size)
    elif cfg.dataset == 'sintel_raw':
        dataset = SINTEL_Prepared(data_dir, num_scales=cfg.num_scales, img_hw=cfg.img_hw, num_iterations=(cfg.num_iterations - cfg.iter_start) * cfg.batch_size)
    elif cfg.dataset == 'kitti_odo':
        dataset = KITTI_Prepared(data_dir, num_scales=cfg.num_scales, img_hw=cfg.img_hw, num_iterations=(cfg.num_iterations - cfg.iter_start) * cfg.batch_size)
    elif cfg.dataset == 'nyuv2':
        dataset = NYU_v2(data_dir, num_scales=cfg.num_scales, img_hw=cfg.img_hw, num_iterations=(cfg.num_iterations - cfg.iter_start) * cfg.batch_size)
    else:
        raise NotImplementedError
    
    dataloader = torch.utils.data.DataLoader(dataset, batch_size=cfg.batch_size, shuffle=True, num_workers=cfg.num_workers, drop_last=False)
    if cfg.dataset == 'kitti_depth' or cfg.dataset == 'kitti_odo' or cfg.dataset == 'sintel_raw':
        gt_flows_2012, noc_masks_2012 = load_gt_flow_kitti(cfg.gt_2012_dir, 'kitti_2012')
        gt_flows_2015, noc_masks_2015 = load_gt_flow_kitti(cfg.gt_2015_dir, 'kitti_2015')
        gt_masks_2015 = load_gt_mask(cfg.gt_2015_dir)
    elif cfg.dataset == 'nyuv2':
        test_images, test_gt_depths = load_nyu_test_data(cfg.nyu_test_dir)

    # training
    print('starting iteration: {}.'.format(cfg.iter_start))
    for iter_, inputs in enumerate(tqdm(dataloader)):
        if (iter_ + 1) % cfg.test_interval == 0 and (not cfg.no_test):
            model.eval()
            if args.multi_gpu:
                model_eval = model.module
            else:
                model_eval = model
            if cfg.dataset == 'kitti_depth' or cfg.dataset == 'kitti_odo' or cfg.dataset == 'sintel_raw':
                if not (cfg.mode == 'depth' or cfg.mode == 'flowposenet'):
                    eval_2012_res = test_kitti_2012(cfg, model_eval, gt_flows_2012, noc_masks_2012)
                    eval_2015_res = test_kitti_2015(cfg, model_eval, gt_flows_2015, noc_masks_2015, gt_masks_2015, depth_save_dir=os.path.join(cfg.model_dir, 'results'))
                    visualizer.add_log_pack({'eval_2012_res': eval_2012_res, 'eval_2015_res': eval_2015_res})
            elif cfg.dataset == 'nyuv2':
                if not cfg.mode == 'flow':
                    eval_nyu_res = test_nyu(cfg, model_eval, test_images, test_gt_depths)
                    visualizer.add_log_pack({'eval_nyu_res': eval_nyu_res})
            visualizer.dump_log(os.path.join(cfg.model_dir, 'log.pkl'))
        model.train()
        iter_ = iter_ + cfg.iter_start
        optimizer.zero_grad()
        inputs = inputs.cuda()
        #inputs = [k.cuda() for k in inputs]
        loss_pack = model(inputs)

        if iter_ % cfg.log_interval == 0:
            visualizer.print_loss(loss_pack, iter_=iter_)

        loss_list = []
        for key in list(loss_pack.keys()):
            loss_list.append((loss_weights_dict[key] * loss_pack[key].mean()).unsqueeze(0))
        loss = torch.cat(loss_list, 0).sum()
        loss.backward()
        optimizer.step()
        if (iter_ + 1) % cfg.save_interval == 0:
            save_model(iter_, cfg.model_dir, 'iter_{}.pth'.format(iter_), model, optimizer)
            save_model(iter_, cfg.model_dir, 'last.pth'.format(iter_), model, optimizer)
    
    if cfg.dataset == 'kitti_depth':
        if cfg.mode == 'depth' or cfg.mode == 'depth_pose':
            eval_depth_res = test_eigen_depth(cfg, model_eval)

if __name__ == '__main__':
    import argparse
    arg_parser = argparse.ArgumentParser(
        description="TrianFlow training pipeline."
    )
    arg_parser.add_argument('-c', '--config_file', default=None, help='config file.')
    arg_parser.add_argument('-g', '--gpu', type=str, default=0, help='gpu id.')
    arg_parser.add_argument('--batch_size', type=int, default=8, help='batch size.')
    arg_parser.add_argument('--iter_start', type=int, default=0, help='starting iteration.')
    arg_parser.add_argument('--lr', type=float, default=0.0001, help='learning rate')
    arg_parser.add_argument('--num_workers', type=int, default=4, help='number of workers.')
    arg_parser.add_argument('--log_interval', type=int, default=100, help='interval for printing loss.')
    arg_parser.add_argument('--test_interval', type=int, default=2000, help='interval for evaluation.')
    arg_parser.add_argument('--save_interval', type=int, default=2000, help='interval for saving models.')
    arg_parser.add_argument('--mode', type=str, default='flow', help='training mode.')
    arg_parser.add_argument('--model_dir', type=str, default=None, help='directory for saving models')
    arg_parser.add_argument('--prepared_save_dir', type=str, default='data_s1', help='directory name for generated training dataset')
    arg_parser.add_argument('--flow_pretrained_model', type=str, default=None, help='directory for loading flow pretrained models')
    arg_parser.add_argument('--depth_pretrained_model', type=str, default=None, help='directory for loading depth pretrained models')
    arg_parser.add_argument('--resume', action='store_true', help='to resume training.')
    arg_parser.add_argument('--multi_gpu', action='store_true', help='to use multiple gpu for training.')
    arg_parser.add_argument('--no_test', action='store_true', help='without evaluation.')
    args = arg_parser.parse_args()
        #args.config_file = 'config/debug.yaml'
    if args.config_file is None:
        raise ValueError('config file needed. -c --config_file.')

    # set model
    if args.model_dir is None:
        args.model_dir = os.path.join('models', os.path.splitext(os.path.split(args.config_file)[1])[0])
    args.model_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), args.model_dir, args.mode)
    if not os.path.exists(args.model_dir):
        os.makedirs(args.model_dir)
    if not os.path.exists(args.config_file):
        raise ValueError('config file not found.')
    with open(args.config_file, 'r') as f:
        cfg = yaml.safe_load(f)
    cfg['img_hw'] = (cfg['img_hw'][0], cfg['img_hw'][1])
    cfg['log_dump_dir'] = os.path.join(args.model_dir, 'log.pkl')
    shutil.copy(args.config_file, args.model_dir)

    # copy attr into cfg
    for attr in dir(args):
        if attr[:2] != '__':
            cfg[attr] = getattr(args, attr)

    # set gpu
    num_gpus = len(args.gpu.split(','))
    if (args.multi_gpu and num_gpus <= 1) or ((not args.multi_gpu) and num_gpus > 1):
        raise ValueError('Error! the number of gpus used in the --gpu argument does not match the argument --multi_gpu.')
    if args.multi_gpu:
        cfg['batch_size'] = cfg['batch_size'] * num_gpus
        cfg['num_iterations'] = int(cfg['num_iterations'] / num_gpus)
    os.environ['CUDA_VISIBLE_DEVICES'] = str(args.gpu)

    class pObject(object):
        def __init__(self):
            pass
    cfg_new = pObject()
    for attr in list(cfg.keys()):
        setattr(cfg_new, attr, cfg[attr])
    with open(os.path.join(args.model_dir, 'config.pkl'), 'wb') as f:
        pickle.dump(cfg_new, f)

    # main function 
    train(cfg_new)