-
Notifications
You must be signed in to change notification settings - Fork 77
/
train.py
178 lines (158 loc) · 6.35 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
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
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
# python imports
import argparse
import os
import time
import datetime
from pprint import pprint
# torch imports
import torch
import torch.nn as nn
import torch.utils.data
# for visualization
from torch.utils.tensorboard import SummaryWriter
# our code
from libs.core import load_config
from libs.datasets import make_dataset, make_data_loader
from libs.modeling import make_meta_arch
from libs.utils import (train_one_epoch, valid_one_epoch, ANETdetection,
save_checkpoint, make_optimizer, make_scheduler,
fix_random_seed, ModelEma)
################################################################################
def main(args):
"""main function that handles training / inference"""
"""1. setup parameters / folders"""
# parse args
args.start_epoch = 0
if os.path.isfile(args.config):
cfg = load_config(args.config)
else:
raise ValueError("Config file does not exist.")
pprint(cfg)
# prep for output folder (based on time stamp)
if not os.path.exists(cfg['output_folder']):
os.mkdir(cfg['output_folder'])
cfg_filename = os.path.basename(args.config).replace('.yaml', '')
if len(args.output) == 0:
ts = datetime.datetime.fromtimestamp(int(time.time()))
ckpt_folder = os.path.join(
cfg['output_folder'], cfg_filename + '_' + str(ts))
else:
ckpt_folder = os.path.join(
cfg['output_folder'], cfg_filename + '_' + str(args.output))
if not os.path.exists(ckpt_folder):
os.mkdir(ckpt_folder)
# tensorboard writer
tb_writer = SummaryWriter(os.path.join(ckpt_folder, 'logs'))
# fix the random seeds (this will fix everything)
rng_generator = fix_random_seed(cfg['init_rand_seed'], include_cuda=True)
# re-scale learning rate / # workers based on number of GPUs
cfg['opt']["learning_rate"] *= len(cfg['devices'])
cfg['loader']['num_workers'] *= len(cfg['devices'])
"""2. create dataset / dataloader"""
train_dataset = make_dataset(
cfg['dataset_name'], True, cfg['train_split'], **cfg['dataset']
)
# update cfg based on dataset attributes (fix to epic-kitchens)
train_db_vars = train_dataset.get_attributes()
cfg['model']['train_cfg']['head_empty_cls'] = train_db_vars['empty_label_ids']
# data loaders
train_loader = make_data_loader(
train_dataset, True, rng_generator, **cfg['loader'])
"""3. create model, optimizer, and scheduler"""
# model
model = make_meta_arch(cfg['model_name'], **cfg['model'])
# not ideal for multi GPU training, ok for now
model = nn.DataParallel(model, device_ids=cfg['devices'])
# optimizer
optimizer = make_optimizer(model, cfg['opt'])
# schedule
num_iters_per_epoch = len(train_loader)
scheduler = make_scheduler(optimizer, cfg['opt'], num_iters_per_epoch)
# enable model EMA
print("Using model EMA ...")
model_ema = ModelEma(model)
"""4. Resume from model / Misc"""
# resume from a checkpoint?
if args.resume:
if os.path.isfile(args.resume):
# load ckpt, reset epoch / best rmse
checkpoint = torch.load(args.resume,
map_location = lambda storage, loc: storage.cuda(
cfg['devices'][0]))
args.start_epoch = checkpoint['epoch']
model.load_state_dict(checkpoint['state_dict'])
model_ema.module.load_state_dict(checkpoint['state_dict_ema'])
# also load the optimizer / scheduler if necessary
optimizer.load_state_dict(checkpoint['optimizer'])
scheduler.load_state_dict(checkpoint['scheduler'])
print("=> loaded checkpoint '{:s}' (epoch {:d}".format(
args.resume, checkpoint['epoch']
))
del checkpoint
else:
print("=> no checkpoint found at '{}'".format(args.resume))
return
# save the current config
with open(os.path.join(ckpt_folder, 'config.txt'), 'w') as fid:
pprint(cfg, stream=fid)
fid.flush()
"""4. training / validation loop"""
print("\nStart training model {:s} ...".format(cfg['model_name']))
# start training
max_epochs = cfg['opt'].get(
'early_stop_epochs',
cfg['opt']['epochs'] + cfg['opt']['warmup_epochs']
)
for epoch in range(args.start_epoch, max_epochs):
# train for one epoch
train_one_epoch(
train_loader,
model,
optimizer,
scheduler,
epoch,
model_ema = model_ema,
clip_grad_l2norm = cfg['train_cfg']['clip_grad_l2norm'],
tb_writer=tb_writer,
print_freq=args.print_freq
)
# save ckpt once in a while
if (
((epoch + 1) == max_epochs) or
((args.ckpt_freq > 0) and ((epoch + 1) % args.ckpt_freq == 0))
):
save_states = {
'epoch': epoch + 1,
'state_dict': model.state_dict(),
'scheduler': scheduler.state_dict(),
'optimizer': optimizer.state_dict(),
}
save_states['state_dict_ema'] = model_ema.module.state_dict()
save_checkpoint(
save_states,
False,
file_folder=ckpt_folder,
file_name='epoch_{:03d}.pth.tar'.format(epoch + 1)
)
# wrap up
tb_writer.close()
print("All done!")
return
################################################################################
if __name__ == '__main__':
"""Entry Point"""
# the arg parser
parser = argparse.ArgumentParser(
description='Train a point-based transformer for action localization')
parser.add_argument('config', metavar='DIR',
help='path to a config file')
parser.add_argument('-p', '--print-freq', default=10, type=int,
help='print frequency (default: 10 iterations)')
parser.add_argument('-c', '--ckpt-freq', default=5, type=int,
help='checkpoint frequency (default: every 5 epochs)')
parser.add_argument('--output', default='', type=str,
help='name of exp folder (default: none)')
parser.add_argument('--resume', default='', type=str, metavar='PATH',
help='path to a checkpoint (default: none)')
args = parser.parse_args()
main(args)