python/main.py at main · tpnam0901/python · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
import logging

logging.root.setLevel(logging.INFO)
logging.basicConfig(level=logging.root.level, format="%(name)s - %(levelname)s - %(message)s")

import argparse
import random

import numpy as np
import torch

import engine
from configs.base import Config, import_config


def arg_parser():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "-cfg",
        "--config",
        type=str,
        default="configs/base.py",
        help="Path to the configuration file.",
    )
    parser.add_argument(
        "-cfg_ckpt",
        "--config_ckpt",
        type=str,
        default="",
        help="Path to the checkpoint .json file to load configuration from.",
    )
    parser.add_argument(
        "-e",
        "--engine",
        choices=["t", "i", "e"],
        required=True,
        help="Engine type to use: TrainEngine (t), InferenceEngine (i), EvaluateEngine (e).",
    )
    return parser.parse_args()


if __name__ == "__main__":

    args = arg_parser()
    cfg: Config = import_config(args.config)

    if args.engine in ["i", "e"]:
        assert (
            args.config_ckpt != ""
        ), "Checkpoint configuration file must be provided for Inference and Evaluation engines."
        cfg.load(args.config_ckpt)

    SEED = cfg.seed
    random.seed(SEED)
    np.random.seed(SEED)
    torch.manual_seed(SEED)
    torch.cuda.manual_seed(SEED)
    torch.cuda.manual_seed_all(SEED)

    if args.engine == "t":
        trainer = engine.TrainEngine(cfg)
        trainer.run()
    elif args.engine == "i":
        inferencer = engine.InferenceEngine(cfg)
        inferencer.run()
    elif args.engine == "e":
        evaluator = engine.EvaluateEngine(cfg)
        evaluator.run()