train.py

import json
import pprint
import datetime
import argparse
from path import Path
from easydict import EasyDict

import basic_train
from logger import init_logger


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('-c', '--config', default='configs/sintel_ft.json')
    parser.add_argument('-e', '--evaluate', action='store_true')
    parser.add_argument('-m', '--model', default=None)
    parser.add_argument('--n_gpu', type=int, default=1)
    args = parser.parse_args()

    with open(args.config) as f:
        cfg = EasyDict(json.load(f))

    if args.evaluate:
        cfg.train.update({
            'epochs': 1,
            'epoch_size': -1,
            'valid_size': 0,
            'workers': 1,
            'val_epoch_size': 1,
        })

    if args.model is not None:
        cfg.train.pretrained_model = args.model
    cfg.train.n_gpu = args.n_gpu

    # store files day by day
    curr_time = datetime.datetime.now().strftime("%y%m%d%H%M%S")
    cfg.save_root = Path('./outputs/checkpoints') / curr_time[:6] / curr_time[6:]
    cfg.save_root.makedirs_p()

    # init logger
    _log = init_logger(log_dir=cfg.save_root, filename=curr_time[6:] + '.log')
    _log.info('=> will save everything to {}'.format(cfg.save_root))

    # show configurations
    cfg_str = pprint.pformat(cfg)
    _log.info('=> configurations \n ' + cfg_str)

    basic_train.main(cfg, _log)