train.py

import os
import argparse
import mlconfig
import shutil
import logging
import numpy as np
import torch

from util import *
from tasks import *


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("--config", type=str)
    return parser.parse_args()

if __name__ == "__main__":
    torch.manual_seed(42)
    np.random.seed(42)
    '''cfg setup'''
    args = parse_args()
    cfg = mlconfig.load(args.config)
    
    V().name = os.path.basename(args.config)
    V().cfg = cfg
    os.makedirs(os.path.join(cfg.output_dir, V().name), exist_ok=True)
    shutil.copy(args.config, os.path.join(cfg.output_dir, V().name))

    '''logger'''
    V().logger = logging.getLogger("")
    V().logger.addHandler(logging.FileHandler(os.path.join(cfg.output_dir, V().name, 'my.log')))
    V().logger.setLevel(logging.INFO)

    '''torch setup'''
    gpu_id = reserve_gpu(None)
    device_id = 0
    V().who_am_I = os.getpid()
    V().device = device_id
    trainer = cfg.trainer()
    
    torch.cuda.set_device(device_id)
    V().info(f'Selecting device {device_id}')

    '''train'''
    trainer.train()