-
Notifications
You must be signed in to change notification settings - Fork 113
/
train.py
222 lines (194 loc) · 9.79 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
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
import argparse
import datetime
import random
import time
from pathlib import Path
import torch
from torch.utils.data import DataLoader, DistributedSampler
from crowd_datasets import build_dataset
from engine import *
from models import build_model
import os
from tensorboardX import SummaryWriter
import warnings
warnings.filterwarnings('ignore')
def get_args_parser():
parser = argparse.ArgumentParser('Set parameters for training P2PNet', add_help=False)
parser.add_argument('--lr', default=1e-4, type=float)
parser.add_argument('--lr_backbone', default=1e-5, type=float)
parser.add_argument('--batch_size', default=8, type=int)
parser.add_argument('--weight_decay', default=1e-4, type=float)
parser.add_argument('--epochs', default=3500, type=int)
parser.add_argument('--lr_drop', default=3500, type=int)
parser.add_argument('--clip_max_norm', default=0.1, type=float,
help='gradient clipping max norm')
# Model parameters
parser.add_argument('--frozen_weights', type=str, default=None,
help="Path to the pretrained model. If set, only the mask head will be trained")
# * Backbone
parser.add_argument('--backbone', default='vgg16_bn', type=str,
help="Name of the convolutional backbone to use")
# * Matcher
parser.add_argument('--set_cost_class', default=1, type=float,
help="Class coefficient in the matching cost")
parser.add_argument('--set_cost_point', default=0.05, type=float,
help="L1 point coefficient in the matching cost")
# * Loss coefficients
parser.add_argument('--point_loss_coef', default=0.0002, type=float)
parser.add_argument('--eos_coef', default=0.5, type=float,
help="Relative classification weight of the no-object class")
parser.add_argument('--row', default=2, type=int,
help="row number of anchor points")
parser.add_argument('--line', default=2, type=int,
help="line number of anchor points")
# dataset parameters
parser.add_argument('--dataset_file', default='SHHA')
parser.add_argument('--data_root', default='./new_public_density_data',
help='path where the dataset is')
parser.add_argument('--output_dir', default='./log',
help='path where to save, empty for no saving')
parser.add_argument('--checkpoints_dir', default='./ckpt',
help='path where to save checkpoints, empty for no saving')
parser.add_argument('--tensorboard_dir', default='./runs',
help='path where to save, empty for no saving')
parser.add_argument('--seed', default=42, type=int)
parser.add_argument('--resume', default='', help='resume from checkpoint')
parser.add_argument('--start_epoch', default=0, type=int, metavar='N',
help='start epoch')
parser.add_argument('--eval', action='store_true')
parser.add_argument('--num_workers', default=8, type=int)
parser.add_argument('--eval_freq', default=5, type=int,
help='frequency of evaluation, default setting is evaluating in every 5 epoch')
parser.add_argument('--gpu_id', default=0, type=int, help='the gpu used for training')
return parser
def main(args):
os.environ["CUDA_VISIBLE_DEVICES"] = '{}'.format(args.gpu_id)
# create the logging file
run_log_name = os.path.join(args.output_dir, 'run_log.txt')
with open(run_log_name, "w") as log_file:
log_file.write('Eval Log %s\n' % time.strftime("%c"))
if args.frozen_weights is not None:
assert args.masks, "Frozen training is meant for segmentation only"
# backup the arguments
print(args)
with open(run_log_name, "a") as log_file:
log_file.write("{}".format(args))
device = torch.device('cuda')
# fix the seed for reproducibility
seed = args.seed + utils.get_rank()
torch.manual_seed(seed)
np.random.seed(seed)
random.seed(seed)
# get the P2PNet model
model, criterion = build_model(args, training=True)
# move to GPU
model.to(device)
criterion.to(device)
model_without_ddp = model
n_parameters = sum(p.numel() for p in model.parameters() if p.requires_grad)
print('number of params:', n_parameters)
# use different optimation params for different parts of the model
param_dicts = [
{"params": [p for n, p in model_without_ddp.named_parameters() if "backbone" not in n and p.requires_grad]},
{
"params": [p for n, p in model_without_ddp.named_parameters() if "backbone" in n and p.requires_grad],
"lr": args.lr_backbone,
},
]
# Adam is used by default
optimizer = torch.optim.Adam(param_dicts, lr=args.lr)
lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, args.lr_drop)
# create the dataset
loading_data = build_dataset(args=args)
# create the training and valiation set
train_set, val_set = loading_data(args.data_root)
# create the sampler used during training
sampler_train = torch.utils.data.RandomSampler(train_set)
sampler_val = torch.utils.data.SequentialSampler(val_set)
batch_sampler_train = torch.utils.data.BatchSampler(
sampler_train, args.batch_size, drop_last=True)
# the dataloader for training
data_loader_train = DataLoader(train_set, batch_sampler=batch_sampler_train,
collate_fn=utils.collate_fn_crowd, num_workers=args.num_workers)
data_loader_val = DataLoader(val_set, 1, sampler=sampler_val,
drop_last=False, collate_fn=utils.collate_fn_crowd, num_workers=args.num_workers)
if args.frozen_weights is not None:
checkpoint = torch.load(args.frozen_weights, map_location='cpu')
model_without_ddp.detr.load_state_dict(checkpoint['model'])
# resume the weights and training state if exists
if args.resume:
checkpoint = torch.load(args.resume, map_location='cpu')
model_without_ddp.load_state_dict(checkpoint['model'])
if not args.eval and 'optimizer' in checkpoint and 'lr_scheduler' in checkpoint and 'epoch' in checkpoint:
optimizer.load_state_dict(checkpoint['optimizer'])
lr_scheduler.load_state_dict(checkpoint['lr_scheduler'])
args.start_epoch = checkpoint['epoch'] + 1
print("Start training")
start_time = time.time()
# save the performance during the training
mae = []
mse = []
# the logger writer
writer = SummaryWriter(args.tensorboard_dir)
step = 0
# training starts here
for epoch in range(args.start_epoch, args.epochs):
t1 = time.time()
stat = train_one_epoch(
model, criterion, data_loader_train, optimizer, device, epoch,
args.clip_max_norm)
# record the training states after every epoch
if writer is not None:
with open(run_log_name, "a") as log_file:
log_file.write("loss/loss@{}: {}".format(epoch, stat['loss']))
log_file.write("loss/loss_ce@{}: {}".format(epoch, stat['loss_ce']))
writer.add_scalar('loss/loss', stat['loss'], epoch)
writer.add_scalar('loss/loss_ce', stat['loss_ce'], epoch)
t2 = time.time()
print('[ep %d][lr %.7f][%.2fs]' % \
(epoch, optimizer.param_groups[0]['lr'], t2 - t1))
with open(run_log_name, "a") as log_file:
log_file.write('[ep %d][lr %.7f][%.2fs]' % (epoch, optimizer.param_groups[0]['lr'], t2 - t1))
# change lr according to the scheduler
lr_scheduler.step()
# save latest weights every epoch
checkpoint_latest_path = os.path.join(args.checkpoints_dir, 'latest.pth')
torch.save({
'model': model_without_ddp.state_dict(),
}, checkpoint_latest_path)
# run evaluation
if epoch % args.eval_freq == 0 and epoch != 0:
t1 = time.time()
result = evaluate_crowd_no_overlap(model, data_loader_val, device)
t2 = time.time()
mae.append(result[0])
mse.append(result[1])
# print the evaluation results
print('=======================================test=======================================')
print("mae:", result[0], "mse:", result[1], "time:", t2 - t1, "best mae:", np.min(mae), )
with open(run_log_name, "a") as log_file:
log_file.write("mae:{}, mse:{}, time:{}, best mae:{}".format(result[0],
result[1], t2 - t1, np.min(mae)))
print('=======================================test=======================================')
# recored the evaluation results
if writer is not None:
with open(run_log_name, "a") as log_file:
log_file.write("metric/mae@{}: {}".format(step, result[0]))
log_file.write("metric/mse@{}: {}".format(step, result[1]))
writer.add_scalar('metric/mae', result[0], step)
writer.add_scalar('metric/mse', result[1], step)
step += 1
# save the best model since begining
if abs(np.min(mae) - result[0]) < 0.01:
checkpoint_best_path = os.path.join(args.checkpoints_dir, 'best_mae.pth')
torch.save({
'model': model_without_ddp.state_dict(),
}, checkpoint_best_path)
# total time for training
total_time = time.time() - start_time
total_time_str = str(datetime.timedelta(seconds=int(total_time)))
print('Training time {}'.format(total_time_str))
if __name__ == '__main__':
parser = argparse.ArgumentParser('P2PNet training and evaluation script', parents=[get_args_parser()])
args = parser.parse_args()
main(args)