-
Notifications
You must be signed in to change notification settings - Fork 1
/
train.py
186 lines (161 loc) · 6.66 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
# This file is part of COAT, and is distributed under the
# OSI-approved BSD 3-Clause License. See top-level LICENSE file or
# https://github.com/Kitware/COAT/blob/master/LICENSE for details.
import argparse
import datetime
import os.path as osp
import time
import torch
import torch.utils.data
from datasets import build_test_loader, build_train_loader
from defaults import get_default_cfg
from engine import evaluate_performance, train_one_epoch
from models.coat import COAT
from utils.utils import mkdir, resume_from_ckpt, save_on_master, set_random_seed
from loss.softmax_loss import SoftmaxLoss #把softmax函数引入
# 计算模型Flops
from ptflops import get_model_complexity_info
def main(args):
cfg = get_default_cfg()
if args.cfg_file:
cfg.merge_from_file(args.cfg_file)
cfg.merge_from_list(args.opts)
cfg.freeze()
device = torch.device(cfg.DEVICE)
if cfg.SEED >= 0:
set_random_seed(cfg.SEED)
print("Creating model...")
model = COAT(cfg)
model.to(device)
# 计算模型的复杂度和大小m和MACs
with torch.cuda.device(0): # todo 1
# net = self.net # Network(self.cfg)
macs, params = get_model_complexity_info(model, (3, 320, 320), as_strings=True,
print_per_layer_stat=True,
verbose=True)
print('{:<30} {:<8}'.format('Computational complexity: ', macs))
print('{:<30} {:<8}'.format('Number of parameters: ', params))
print("Loading data...")
train_loader = build_train_loader(cfg)
gallery_loader, query_loader = build_test_loader(cfg)
softmax_criterion_s2 = None #先初始化定义 这些损失为空
softmax_criterion_s3 = None
if cfg.MODEL.LOSS.USE_SOFTMAX:
softmax_criterion_s2 = SoftmaxLoss(cfg) #引入的softamx函数 第二个阶段
softmax_criterion_s3 = SoftmaxLoss(cfg) #引入的softmax第三个阶段
softmax_criterion_s2.to(device)
softmax_criterion_s3.to(device)
if args.eval:
assert args.ckpt, "--ckpt must be specified when --eval enabled"
resume_from_ckpt(args.ckpt, model)
evaluate_performance(
model,
gallery_loader,
query_loader,
device,
use_gt=cfg.EVAL_USE_GT,
use_cache=cfg.EVAL_USE_CACHE,
use_cbgm=cfg.EVAL_USE_CBGM,
gallery_size=cfg.EVAL_GALLERY_SIZE,
)
exit(0)
params = [p for p in model.parameters() if p.requires_grad]
if cfg.MODEL.LOSS.USE_SOFTMAX:
params_softmax_s2 = [p for p in softmax_criterion_s2.parameters() if p.requires_grad]
params_softmax_s3 = [p for p in softmax_criterion_s3.parameters() if p.requires_grad]
params.extend(params_softmax_s2)
params.extend(params_softmax_s3)
optimizer = torch.optim.SGD(
params,
lr=cfg.SOLVER.BASE_LR,
momentum=cfg.SOLVER.SGD_MOMENTUM,
weight_decay=cfg.SOLVER.WEIGHT_DECAY,
)
lr_scheduler = torch.optim.lr_scheduler.MultiStepLR(
optimizer, milestones=cfg.SOLVER.LR_DECAY_MILESTONES, gamma=cfg.SOLVER.GAMMA
)
start_epoch = 0
if args.resume:
assert args.ckpt, "--ckpt must be specified when --resume enabled"
start_epoch = resume_from_ckpt(args.ckpt, model, optimizer, lr_scheduler) + 1
print("Creating output folder...")
output_dir = cfg.OUTPUT_DIR
mkdir(output_dir)
path = osp.join(output_dir, "config.yaml")
with open(path, "w") as f:
f.write(cfg.dump())
print(f"Full config is saved to {path}")
tfboard = None
if cfg.TF_BOARD:
from torch.utils.tensorboard import SummaryWriter
tf_log_path = osp.join(output_dir, "tf_log")
mkdir(tf_log_path)
tfboard = SummaryWriter(log_dir=tf_log_path)
print(f"TensorBoard files are saved to {tf_log_path}")
print("Start training...")
start_time = time.time()
for epoch in range(start_epoch, cfg.SOLVER.MAX_EPOCHS):
train_one_epoch(cfg, model, optimizer, train_loader, device, epoch, tfboard, softmax_criterion_s2, softmax_criterion_s3)
lr_scheduler.step()
if (epoch + 1) % cfg.EVAL_PERIOD == 0 or epoch == cfg.SOLVER.MAX_EPOCHS - 1: #EVAL_PERIOD=1
evaluate_performance(
model,
gallery_loader,
query_loader,
device,
use_gt=cfg.EVAL_USE_GT,
use_cache=cfg.EVAL_USE_CACHE,
use_cbgm=cfg.EVAL_USE_CBGM,
)
if (epoch + 1) % cfg.CKPT_PERIOD == 0 or epoch == cfg.SOLVER.MAX_EPOCHS - 1: #CKPT_PERIOD=1
save_on_master(
{
"model": model.state_dict(),
"optimizer": optimizer.state_dict(),
"lr_scheduler": lr_scheduler.state_dict(),
"epoch": epoch,
},
osp.join(output_dir, f"epoch_{epoch}.pth"),
)
# # only save the last three checkpoints
# if epoch >= cfg.SOLVER.MAX_EPOCHS - 3:
# save_on_master(
# {
# "model": model.state_dict(),
# "optimizer": optimizer.state_dict(),
# "lr_scheduler": lr_scheduler.state_dict(),
# "epoch": epoch,
# },
# osp.join(output_dir, f"epoch_{epoch}.pth"),
# )
# # evaluate the current checkpoint
# evaluate_performance(
# model,
# gallery_loader,
# query_loader,
# device,
# use_gt=cfg.EVAL_USE_GT,
# use_cache=cfg.EVAL_USE_CACHE,
# use_cbgm=cfg.EVAL_USE_CBGM,
# gallery_size=cfg.EVAL_GALLERY_SIZE,
# )
if tfboard:
tfboard.close()
total_time = time.time() - start_time
total_time_str = str(datetime.timedelta(seconds=int(total_time)))
print(f"Total training time {total_time_str}")
if __name__ == "__main__":
parser = argparse.ArgumentParser(description="Train a person search network.")
parser.add_argument("--cfg", dest="cfg_file", help="Path to configuration file.")
parser.add_argument(
"--eval", action="store_true", help="Evaluate the performance of a given checkpoint."
)
parser.add_argument(
"--resume", action="store_true", help="Resume from the specified checkpoint."
)
parser.add_argument("--ckpt", help="Path to checkpoint to resume or evaluate.")
parser.add_argument(
"opts", nargs=argparse.REMAINDER, help="Modify config options using the command-line"
)
args = parser.parse_args()
main(args)