SilVar/train.py at main · Hanhpt23/SilVar · 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
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
"""
 Copyright (c) 2022, salesforce.com, inc.
 All rights reserved.
 SPDX-License-Identifier: BSD-3-Clause
 For full license text, see the LICENSE_Lavis file in the repo root or https://opensource.org/licenses/BSD-3-Clause
"""

import argparse
import os
import glob
import random
import yaml

import numpy as np
import torch
import torch.backends.cudnn as cudnn
import wandb

import silvar.tasks as tasks
from silvar.common.config import Config
from silvar.common.dist_utils import get_rank, init_distributed_mode
from silvar.common.logger import setup_logger
from silvar.common.optims import (
    LinearWarmupCosineLRScheduler,
    LinearWarmupStepLRScheduler,
)
from silvar.common.registry import registry
from silvar.common.utils import now

# imports modules for registration
from silvar.models import *
from evaluate import *
from silvar.processors import *
from silvar.runners import *
from silvar.tasks import *

def list_of_str(arg):
    return list(map(str, arg.split(',')))

def parse_args():
    parser = argparse.ArgumentParser(description="Training")

    parser.add_argument("--cfg-path", required=True, help="path to train configuration file.")
    parser.add_argument("--cfg-eval-path", required=False, help="path to evaluation configuration file.")
    parser.add_argument("--eval-dataset", type=list_of_str, default='val_vindrcxr', help="dataset to evaluate")
    parser.add_argument(
        "--options",
        nargs="+",
        help="override some settings in the used config, the key-value pair "
        "in xxx=yyy format will be merged into config file (deprecate), "
        "change to --cfg-options instead.",
    )
    args = parser.parse_args()

    return args


def setup_seeds(config):
    seed = config.run_cfg.seed + get_rank()

    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)

    cudnn.benchmark = False
    cudnn.deterministic = True


def get_runner_class(cfg):
    """
    Get runner class from config. Default to epoch-based runner.
    """
    runner_cls = registry.get_runner_class(cfg.run_cfg.get("runner", "runner_base"))

    return runner_cls


def main():

    # set before init_distributed_mode() to ensure the same job_id shared across all ranks.
    job_id = now()
    args = parse_args()
    cfg = Config(args)

    print(cfg)

    init_distributed_mode(cfg.run_cfg)
    setup_seeds(cfg)

    # set after init_distributed_mode() to only log on master.
    setup_logger()
    cfg.pretty_print()

    task = tasks.setup_task(cfg)
    datasets = task.build_datasets(cfg)
    model = task.build_model(cfg)

    if cfg.run_cfg.wandb_log:
        wandb.login(key=cfg.run_cfg.wandb_token)
        wandb.init(project="ars2text", name=cfg.run_cfg.job_name)
        wandb.watch(model)

    runner = get_runner_class(cfg)(
        cfg=cfg, job_id=job_id, task=task, model=model, datasets=datasets
    )
    runner.train()

    if hasattr(args, 'cfg_eval_path'):
        args.cfg_path = args.cfg_eval_path

        model_path = "silvar/{}/{}".format(cfg.run_cfg.output_dir, job_id)
        ckpt_paths = glob.glob(os.path.join(model_path, "*.pth"))
        ckpt_names = [os.path.basename(ckp_path) for ckp_path in ckpt_paths]
        last_ckpt_name = sorted(ckpt_names, key=lambda x: int(x.split(".")[0].split("_")[-1]))[-1]
        last_ckpt_path = os.path.join(model_path, last_ckpt_name)

        with open(args.cfg_path) as f:
            eval_cfg = yaml.load(f, Loader=yaml.FullLoader)
            eval_cfg["model"]["ckpt"] = last_ckpt_path

        with open(args.cfg_path, "w") as f:
            yaml.dump(
                eval_cfg, stream=f, default_flow_style=False, sort_keys=False
            )

        print("Evaluating...........")
        evaluate(args)
        print("Done!")

if __name__ == "__main__":
    main()