-
Notifications
You must be signed in to change notification settings - Fork 7
/
train.py
84 lines (61 loc) · 2.68 KB
/
train.py
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
"""
DEIM: DETR with Improved Matching for Fast Convergence
Copyright (c) 2024 The DEIM Authors. All Rights Reserved.
---------------------------------------------------------------------------------
Modified from RT-DETR (https://github.com/lyuwenyu/RT-DETR)
Copyright (c) 2023 lyuwenyu. All Rights Reserved.
"""
import os
import sys
sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), '..'))
import argparse
from engine.misc import dist_utils
from engine.core import YAMLConfig, yaml_utils
from engine.solver import TASKS
debug=False
if debug:
import torch
def custom_repr(self):
return f'{{Tensor:{tuple(self.shape)}}} {original_repr(self)}'
original_repr = torch.Tensor.__repr__
torch.Tensor.__repr__ = custom_repr
def main(args, ) -> None:
"""main
"""
dist_utils.setup_distributed(args.print_rank, args.print_method, seed=args.seed)
assert not all([args.tuning, args.resume]), \
'Only support from_scrach or resume or tuning at one time'
update_dict = yaml_utils.parse_cli(args.update)
update_dict.update({k: v for k, v in args.__dict__.items() \
if k not in ['update', ] and v is not None})
cfg = YAMLConfig(args.config, **update_dict)
if args.resume or args.tuning:
if 'HGNetv2' in cfg.yaml_cfg:
cfg.yaml_cfg['HGNetv2']['pretrained'] = False
print('cfg: ', cfg.__dict__)
solver = TASKS[cfg.yaml_cfg['task']](cfg)
if args.test_only:
solver.val()
else:
solver.fit()
dist_utils.cleanup()
if __name__ == '__main__':
parser = argparse.ArgumentParser()
# priority 0
parser.add_argument('-c', '--config', type=str, required=True)
parser.add_argument('-r', '--resume', type=str, help='resume from checkpoint')
parser.add_argument('-t', '--tuning', type=str, help='tuning from checkpoint')
parser.add_argument('-d', '--device', type=str, help='device',)
parser.add_argument('--seed', type=int, help='exp reproducibility')
parser.add_argument('--use-amp', action='store_true', help='auto mixed precision training')
parser.add_argument('--output-dir', type=str, help='output directoy')
parser.add_argument('--summary-dir', type=str, help='tensorboard summry')
parser.add_argument('--test-only', action='store_true', default=False,)
# priority 1
parser.add_argument('-u', '--update', nargs='+', help='update yaml config')
# env
parser.add_argument('--print-method', type=str, default='builtin', help='print method')
parser.add_argument('--print-rank', type=int, default=0, help='print rank id')
parser.add_argument('--local-rank', type=int, help='local rank id')
args = parser.parse_args()
main(args)