-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathtrain.py
191 lines (162 loc) · 8.55 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
179
180
181
182
183
184
185
186
187
188
189
190
191
import os
import sys
import time
import subprocess
from typing import Any, Dict, List, Tuple, Union
from datetime import datetime
import argparse
import faulthandler
from tqdm import tqdm
#
import torch
import torch.multiprocessing
from torch.utils.data import DataLoader
#
from loader import Loader
from utils.logger import Logger
from utils.utils import AverageMeterForDict
from utils.utils import save_ckpt, set_seed
from utils import time_utils as tu
from utils import plot_utils as pu
import matplotlib.pyplot as plt
def parse_arguments() -> Any:
"""Arguments for running the baseline.
Returns:
parsed arguments
"""
parser = argparse.ArgumentParser()
parser.add_argument("--mode", default="train", type=str, help="Mode, train/val/test")
parser.add_argument("--features_dir", required=True, default="/private/wangchen/instance_model/instance_model_data_simpl/", type=str, help="Path to the dataset")
parser.add_argument("--train_batch_size", type=int, default=16, help="Training batch size")
parser.add_argument("--val_batch_size", type=int, default=16, help="Val batch size")
parser.add_argument("--train_epoches", type=int, default=10, help="Number of epoches for training")
parser.add_argument("--val_interval", type=int, default=5, help="Validation intervals")
parser.add_argument("--data_aug", action="store_true", help="Enable data augmentation")
parser.add_argument("--seed", type=int, default=42, help="Random seed")
parser.add_argument("--use_cuda", action="store_true", help="Use CUDA for acceleration")
parser.add_argument("--logger_writer", action="store_true", help="Enable tensorboard")
parser.add_argument("--adv_cfg_path", required=True, default="", type=str)
parser.add_argument("--rank_metric", required=False, type=str, default="brier_fde_k", help="Ranking metric")
parser.add_argument("--resume", action="store_true", help="Resume training")
parser.add_argument("--no_pbar", action="store_true", help="Hide progress bar")
parser.add_argument("--model_path", required=False, type=str, help="path to the saved model")
return parser.parse_args()
def main():
args = parse_arguments()
faulthandler.enable()
start_time = time.time()
set_seed(args.seed)
if args.use_cuda and torch.cuda.is_available():
device = torch.device("cuda", 0)
else:
device = torch.device('cpu')
date_str = datetime.now().strftime("%Y%m%d-%H%M%S")
log_dir = "log/" + date_str
logger = Logger(date_str=date_str, log_dir=log_dir, enable_flags={'writer': args.logger_writer})
# log basic info
logger.log_basics(args=args, datetime=date_str)
loader = Loader(args, device, is_ddp=False)
if args.resume:
logger.print('[Resume] Loading state_dict from {}'.format(args.model_path))
loader.set_resmue(args.model_path)
(train_set, val_set), net, loss_fn, optimizer, evaluator = loader.load()
dl_train = DataLoader(train_set,
batch_size=args.train_batch_size,
shuffle=True,
num_workers=8,
collate_fn=train_set.collate_fn,
drop_last=True,
pin_memory=True)
dl_val = DataLoader(val_set,
batch_size=args.val_batch_size,
shuffle=False,
num_workers=8,
collate_fn=val_set.collate_fn,
drop_last=True,
pin_memory=True)
niter = 0
best_metric = 1e6
rank_metric = args.rank_metric
net_name = loader.network_name()
print("d:begin_train")
for epoch in range(args.train_epoches):
logger.print('\nEpoch {}'.format(epoch))
torch.cuda.empty_cache()
torch.cuda.reset_peak_memory_stats()
# * Train
epoch_start = time.time()
train_loss_meter = AverageMeterForDict()
train_eval_meter = AverageMeterForDict()
net.train()
for i, data in enumerate(tqdm(dl_train, disable=args.no_pbar, ncols=80)):
data_in = net.pre_process(data)
# actors：(all_actor_Num of batch, 3, 20) 3代表xy_delta(也即xy方向上的v)和pad 20是obs_len
# actor_idcs :list len = size of batch 因为actors包含了batch所有obejct，因此idcs为了指明哪个属于一个sample
# 同理lane:(all_lane_Num, 10, 10) 也即(all seg_num of batch, 每个seg采10个点, 堆叠后的维度10代表node_ctrs/turn/control等信息)
# rpe ,list len = bs, dict, scene = (5, x, x) x= n_agent + n_lane
actors, actor_idcs, lanes, lane_idcs, rpe = data_in
# print("d:",data_in)
out = net(data_in)
res_cls, res_reg, res_aux = out # res_cls(bs, agent_num, 6) res_reg(bs, agent_num, 6, 30, 2)
loss_out = loss_fn(out, data)
post_out = net.post_process(out, data,draw=False)
eval_out = evaluator.evaluate(post_out, data)
optimizer.zero_grad()
loss_out['loss'].backward()
lr = optimizer.step()
train_loss_meter.update(loss_out)
train_eval_meter.update(eval_out)
niter += args.train_batch_size
logger.add_dict(loss_out, niter, prefix='train/')
# print('epoch: {}, lr: {}'.format(epoch, lr))
optimizer.step_scheduler()
max_memory = torch.cuda.max_memory_allocated(device=device) // 2 ** 20
loss_avg = train_loss_meter.metrics['loss'].avg
logger.print('[Training] Avg. loss: {:.6}, time cost: {:.3} mins, lr: {:.3}, peak mem: {} MB'.
format(loss_avg, (time.time() - epoch_start) / 60.0, lr, max_memory))
logger.print('-- ' + train_eval_meter.get_info())
logger.add_scalar('train/lr', lr, it=epoch)
logger.add_scalar('train/max_mem', max_memory, it=epoch)
for key, elem in train_eval_meter.metrics.items():
logger.add_scalar(title='train/{}'.format(key), value=elem.avg, it=epoch)
if ((epoch + 1) % args.val_interval == 0) or epoch > int(args.train_epoches / 2):
# * Validation
with torch.no_grad():
val_start = time.time()
val_loss_meter = AverageMeterForDict()
val_eval_meter = AverageMeterForDict()
net.eval()
for i, data in enumerate(tqdm(dl_val, disable=args.no_pbar, ncols=80)):
data_in = net.pre_process(data)
out = net(data_in)
loss_out = loss_fn(out, data)
post_out = net.post_process(out,data,draw= False)
eval_out = evaluator.evaluate(post_out, data)
val_loss_meter.update(loss_out)
val_eval_meter.update(eval_out)
logger.print('[Validation] Avg. loss: {:.6}, time cost: {:.3} mins'.format(
val_loss_meter.metrics['loss'].avg, (time.time() - val_start) / 60.0))
logger.print('-- ' + val_eval_meter.get_info())
for key, elem in val_loss_meter.metrics.items():
logger.add_scalar(title='val/{}'.format(key), value=elem.avg, it=epoch)
for key, elem in val_eval_meter.metrics.items():
logger.add_scalar(title='val/{}'.format(key), value=elem.avg, it=epoch)
if (epoch >= args.train_epoches / 2):
if val_eval_meter.metrics[rank_metric].avg < best_metric:
model_name = date_str + '_{}_best.tar'.format(net_name)
save_ckpt(net, optimizer, epoch, 'saved_models/', model_name)
best_metric = val_eval_meter.metrics[rank_metric].avg
print('Save the model: {}, {}: {:.4}, epoch: {}'.format(
model_name, rank_metric, best_metric, epoch))
if int(100 * epoch / args.train_epoches) in [20, 40, 60, 80]:
model_name = date_str + '_{}_ckpt_epoch{}.tar'.format(net_name, epoch)
save_ckpt(net, optimizer, epoch, 'saved_models/', model_name)
logger.print('Save the model to {}'.format('saved_models/' + model_name))
logger.print("\nTraining completed in {:.2f} mins".format((time.time() - start_time) / 60.0))
# save trained model
model_name = date_str + '_{}_epoch{}.tar'.format(net_name, args.train_epoches)
save_ckpt(net, optimizer, epoch, 'saved_models/', model_name)
print('Save the model to {}'.format('saved_models/' + model_name))
print('\nExit...\n')
if __name__ == "__main__":
main()