main.py

#!/usr/bin/env python3

__doc__ = ""
import sys, os

sys.path.append(os.path.join(os.path.dirname(__file__), "runners"))

import argparse
import traceback
import time
import shutil
import logging
import yaml
import torch
import numpy as np
from runners import *


def parse_args_and_config():
    parser = argparse.ArgumentParser(description=globals()["__doc__"])

    parser.add_argument(
        "--runner", type=str, default="ToyRunner", help="The runner to execute"
    )
    parser.add_argument(
        "--config", type=str, default="toy.yml", help="Path to the config file"
    )
    parser.add_argument("--seed", type=int, default=0, help="Random seed")
    parser.add_argument(
        "--run",
        type=str,
        default="storage",
        help="Path for saving running related data.",
    )
    parser.add_argument(
        "--doc", type=str, default="0", help="A string for documentation purpose"
    )
    parser.add_argument(
        "--comment", type=str, default="", help="A string for experiment comment"
    )
    parser.add_argument(
        "--verbose",
        type=str,
        default="info",
        help="Verbose level: info | debug | warning | critical",
    )
    parser.add_argument("--test", action="store_true", help="Whether to test the model")
    parser.add_argument(
        "--resume_training", action="store_true", help="Whether to resume training"
    )
    parser.add_argument(
        "-o",
        "--image_folder",
        type=str,
        default="images",
        help="The directory of image outputs",
    )

    args = parser.parse_args()
    run_id = str(os.getpid())
    run_time = time.strftime("%Y-%b-%d-%H-%M-%S")
    # args.doc = '_'.join([args.doc, run_id, run_time])
    args.log = os.path.join(args.run, "logs", args.doc)

    # parse config file
    if not args.test:
        with open(os.path.join("configs", args.config), "r") as f:
            config = yaml.safe_load(f)
        new_config = dict2namespace(config)
    else:
        # with open(os.path.join(args.log, 'config.yml'), 'r') as f:
        with open(os.path.join("configs", args.config), "r") as f:
            config = yaml.safe_load(f)
        # new_config = config
        new_config = dict2namespace(config)

    if not args.test:
        if not args.resume_training:
            if os.path.exists(args.log):
                shutil.rmtree(args.log)
            os.makedirs(args.log)

        with open(os.path.join(args.log, "config.yml"), "w") as f:
            yaml.dump(new_config, f, default_flow_style=False)

        # setup logger
        level = getattr(logging, args.verbose.upper(), None)
        if not isinstance(level, int):
            raise ValueError("level {} not supported".format(args.verbose))

        handler1 = logging.StreamHandler()
        handler2 = logging.FileHandler(os.path.join(args.log, "stdout.txt"))
        formatter = logging.Formatter(
            "%(levelname)-8s - %(filename)-15s - %(asctime)s - %(message)s"
        )
        handler1.setFormatter(formatter)
        handler2.setFormatter(formatter)
        logger = logging.getLogger()
        logger.addHandler(handler1)
        logger.addHandler(handler2)
        logger.setLevel(level)

    else:
        level = getattr(logging, args.verbose.upper(), None)
        if not isinstance(level, int):
            raise ValueError("level {} not supported".format(args.verbose))

        handler1 = logging.StreamHandler()
        formatter = logging.Formatter(
            "%(levelname)-8s - %(filename)-15s - %(asctime)s - %(message)s"
        )
        handler1.setFormatter(formatter)
        logger = logging.getLogger()
        logger.addHandler(handler1)
        logger.setLevel(level)

    if not os.path.exists(new_config.storage_path):
        os.makedirs(new_config.storage_path)
    # add device
    device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
    logging.info("Using device: {}".format(device))
    new_config.device = device

    # set random seed
    torch.manual_seed(args.seed)
    np.random.seed(args.seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(args.seed)

    torch.backends.cudnn.benchmark = True

    return args, new_config


def dict2namespace(config):
    namespace = argparse.Namespace()
    for key, value in config.items():
        if isinstance(value, dict):
            new_value = dict2namespace(value)
        else:
            new_value = value
        setattr(namespace, key, new_value)
    return namespace


def main():
    args, config = parse_args_and_config()
    logging.info("Writing log file to {}".format(args.log))
    logging.info("Exp instance id = {}".format(os.getpid()))
    logging.info("Exp comment = {}".format(args.comment))
    logging.info("Config =")
    logging.info(">" * 80)
    logging.info(config)
    logging.info("<" * 80)

    base_path = os.path.dirname(__file__)
    config.base_path = base_path
    try:
        runner = eval(args.runner)(args, config)
        if not args.test:
            runner.train()
        else:
            runner.test()
    except:
        logging.warning(traceback.format_exc())

    return 0


if __name__ == "__main__":
    sys.exit(main())