-
Notifications
You must be signed in to change notification settings - Fork 0
/
train_coco.py
208 lines (177 loc) · 7.17 KB
/
train_coco.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
import torch
import torch.multiprocessing as mp
import torch.distributed as dist
from torchvision import transforms
import random
from functools import partial
from easydict import EasyDict as edict
from albumentations import (
Compose, HorizontalFlip, ShiftScaleRotate, PadIfNeeded, RandomCrop,
RGBShift, RandomBrightness, RandomContrast
)
from adaptis.engine.trainer import AdaptISTrainer
from adaptis.model.cityscapes.models import get_cityscapes_model
from adaptis.model.losses import NormalizedFocalLossSigmoid, NormalizedFocalLossSoftmax, AdaptISProposalsLossIoU
from adaptis.model.metrics import AdaptiveIoU
from adaptis.data.coco import COCODataset
from adaptis.utils import log
from adaptis.model import initializer
from adaptis.utils.exp import init_experiment
def add_exp_args(parser):
parser.add_argument('--dataset-path', type=str, help='Path to the dataset')
return parser
def init_model(args):
model_cfg = edict()
model_cfg.syncbn = True
model_cfg.crop_size = (544, 544)
model_cfg.input_normalization = {
'mean': [.485, .456, .406],
'std': [.229, .224, .225]
}
model_cfg.input_transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize(model_cfg.input_normalization['mean'],
model_cfg.input_normalization['std']),
])
if args.ngpus > 1 and model_cfg.syncbn:
norm_layer = torch.nn.SyncBatchNorm
else:
norm_layer = torch.nn.BatchNorm2d
model = get_cityscapes_model(num_classes=133, norm_layer=norm_layer,
backbone='resnet50')
model.apply(initializer.XavierGluon(rnd_type='gaussian', magnitude=2.0))
model.backbone.load_pretrained_weights()
return model, model_cfg
def scale_func(image_shape):
return random.uniform(350, 700) / min(*image_shape[:2])
def train(model, model_cfg, args, train_proposals, start_epoch=0,
rank=None,
mp_distributed=False,
world_size=None):
args.val_batch_size = args.batch_size
args.input_normalization = model_cfg.input_normalization
crop_size = model_cfg.crop_size
loss_cfg = edict()
loss_cfg.instance_loss = NormalizedFocalLossSigmoid(alpha=0.25, gamma=2)
loss_cfg.instance_loss_weight = 1.0 if not train_proposals else 0.0
if not train_proposals:
num_epochs = 20
num_points = 8
loss_cfg.segmentation_loss = NormalizedFocalLossSoftmax(ignore_label=-1, gamma=1)
loss_cfg.segmentation_loss_weight = 0.75
else:
num_epochs = 2 # rt: I set it
num_points = 48
loss_cfg.proposals_loss = AdaptISProposalsLossIoU(args.batch_size)
loss_cfg.proposals_loss_weight = 1.0
train_augmentator = Compose([
HorizontalFlip(),
# ShiftScaleRotate(shift_limit=0.03, scale_limit=0,
# rotate_limit=(-3, 3), border_mode=0, p=0.75),
PadIfNeeded(min_height=crop_size[0], min_width=crop_size[1], border_mode=0),
RandomCrop(*crop_size),
# RandomBrightness(limit=(-0.25, 0.25), p=0.75),
# RandomContrast(limit=(-0.15, 0.4), p=0.75),
# RGBShift(r_shift_limit=10, g_shift_limit=10, b_shift_limit=10, p=0.75)
], p=1.0)
val_augmentator = Compose([
PadIfNeeded(min_height=crop_size[0], min_width=crop_size[1], border_mode=0),
RandomCrop(*crop_size)
], p=1.0)
trainset = COCODataset(
args.dataset_path,
split='train',
num_points=num_points,
augmentator=train_augmentator,
with_segmentation=True,
points_from_one_object=train_proposals,
input_transform=model_cfg.input_transform,
min_object_area=80,
sample_ignore_object_prob=0.025,
keep_background_prob=0.05,
get_image_scale=scale_func,
use_jpeg=False
)
valset = COCODataset(
args.dataset_path,
split='val',
augmentator=val_augmentator,
num_points=num_points,
with_segmentation=True,
points_from_one_object=train_proposals,
input_transform=model_cfg.input_transform,
min_object_area=80,
keep_background_prob=0.01,
get_image_scale=scale_func,
use_jpeg=False
)
if not train_proposals:
optimizer_params = {
'lr': 1e-4,
'betas': (0.9, 0.999), 'eps': 1e-8, 'weight_decay': 1e-4
}
lr_scheduler = lambda *args, **kwargs: \
partial(torch.optim.lr_scheduler.MultiStepLR,
milestones=[kwargs['T_max'] / 20 * 15, kwargs['T_max'] / 20 * 19],
gamma=0.1,
last_epoch=-1)(*args, **{k:v for k,v in kwargs.items() if k != 'T_max'})
for p in model.parameters():
if hasattr(p, 'lr_mult'):
assert p.lr_mult == 0.1
p.lr_mult = 0.5
else:
optimizer_params = {
'lr': 5e-4,
'betas': (0.9, 0.999), 'eps': 1e-8
}
lr_scheduler = partial(torch.optim.lr_scheduler.CosineAnnealingLR,
last_epoch=-1)
trainer = AdaptISTrainer(args, model, model_cfg, loss_cfg,
trainset, valset,
num_epochs=num_epochs,
optimizer='adam',
optimizer_params=optimizer_params,
lr_scheduler=lr_scheduler,
checkpoint_interval=40 if not train_proposals else 2,
image_dump_interval=100 if not train_proposals else -1,
train_proposals=train_proposals,
metrics=[AdaptiveIoU()],
rank=rank,
mp_distributed=mp_distributed,
world_size=world_size)
if args.resume:
start_epoch = trainer.load_last_checkpoint() + 1
log.logger.info(f'Starting Epoch: {start_epoch}')
log.logger.info(f'Total Epochs: {num_epochs}')
for epoch in range(start_epoch, num_epochs):
trainer.training(epoch)
# trainer.validation(epoch)
def main(rank, ngpus, args, port):
distributed = rank is not None # and not debug
if distributed: # multiprocess distributed training
dist.init_process_group(
world_size=ngpus, rank=rank,
backend='nccl', init_method=f'tcp://127.0.0.1:{port}',
)
torch.cuda.set_device(rank)
model, model_cfg = init_model(args)
train(model, model_cfg, args, train_proposals=False,
start_epoch=args.start_epoch,
rank=rank,
mp_distributed=distributed,
world_size=ngpus)
model.add_proposals_head()
train(model, model_cfg, args, train_proposals=True,
start_epoch=args.start_epoch,
rank=rank,
mp_distributed=distributed,
world_size=ngpus)
if __name__ == '__main__':
args = init_experiment('coco', add_exp_args, script_path=__file__)
ngpus = torch.cuda.device_count()
port = random.randint(10000, 20000)
argv = (ngpus, args, port)
if args.dist:
mp.spawn(main, nprocs=ngpus, args=argv)
else:
main(None, *argv)