main.py

import os
import torch
import argparse
from util.conf import Config, load_config
from util.log import init_logger
from util.utils import seed_it
from trainer import Trainer

os.makedirs('logs', exist_ok=True)

class Coach:
    def __init__(self, config: Config, **kwargs):
        self.config = config
        self.device = torch.device(f'cuda:{config.base.gpu}' if torch.cuda.is_available() else 'cpu')
        self.trainer = Trainer(config, **kwargs)

    def run(self):
        self.trainer.train()
    
    def eval(self):
        self.trainer.eval()


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Model Configs')
    parser.add_argument('--config', '-c', default='conf/mmst.toml', type=str, help='config file path')
    parser.add_argument('--eval', '-e', default=None, type=str, help='model path for evaluation')
    parser.add_argument('--device', '-d', default='cpu', type=str, help='cpu or cuda device')
    args = parser.parse_args()

    kwargs = {
        'model_path': args.eval,
        'device': args.device
    }

    config = load_config(args.config)
    seed_it(config.base.seed)

    if args.eval is not None:
        log = init_logger('main', f'eval_{config.base.name}')
        kwargs = {'model_path': args.eval, 'device': args.device}
        coach = Coach(config, **kwargs)
        coach.eval()
    else:
        log = init_logger('main', config.base.name)
        coach = Coach(config)

        for section, options in config.__dict__.items():
            if isinstance(options, dict):
                log.info(f"[{section}]")
                for key, value in options.items():
                    log.info(f"  {key}: {value}")
            else:
                log.info(f"{section}: {options}")
        
        coach.run()