-
Notifications
You must be signed in to change notification settings - Fork 50
/
training.py
179 lines (147 loc) · 7.17 KB
/
training.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
"""
Training the model
Extended from original implementation of PANet by Wang et al.
"""
import os
import shutil
import torch
import torch.nn as nn
import torch.optim
from torch.utils.data import DataLoader
from torch.optim.lr_scheduler import MultiStepLR
import torch.backends.cudnn as cudnn
import numpy as np
from models.grid_proto_fewshot import FewShotSeg
from dataloaders.dev_customized_med import med_fewshot
from dataloaders.GenericSuperDatasetv2 import SuperpixelDataset
from dataloaders.dataset_utils import DATASET_INFO
import dataloaders.augutils as myaug
from util.utils import set_seed, t2n, to01, compose_wt_simple
from util.metric import Metric
from config_ssl_upload import ex
import tqdm
# config pre-trained model caching path
os.environ['TORCH_HOME'] = "./pretrained_model"
@ex.automain
def main(_run, _config, _log):
if _run.observers:
os.makedirs(f'{_run.observers[0].dir}/snapshots', exist_ok=True)
for source_file, _ in _run.experiment_info['sources']:
os.makedirs(os.path.dirname(f'{_run.observers[0].dir}/source/{source_file}'),
exist_ok=True)
_run.observers[0].save_file(source_file, f'source/{source_file}')
shutil.rmtree(f'{_run.observers[0].basedir}/_sources')
set_seed(_config['seed'])
cudnn.enabled = True
cudnn.benchmark = True
torch.cuda.set_device(device=_config['gpu_id'])
torch.set_num_threads(1)
_log.info('###### Create model ######')
model = FewShotSeg(pretrained_path=None, cfg=_config['model'])
model = model.cuda()
model.train()
_log.info('###### Load data ######')
### Training set
data_name = _config['dataset']
if data_name == 'SABS_Superpix':
baseset_name = 'SABS'
elif data_name == 'C0_Superpix':
raise NotImplementedError
baseset_name = 'C0'
elif data_name == 'CHAOST2_Superpix':
baseset_name = 'CHAOST2'
else:
raise ValueError(f'Dataset: {data_name} not found')
### Transforms for data augmentation
tr_transforms = myaug.transform_with_label({'aug': myaug.augs[_config['which_aug']]})
assert _config['scan_per_load'] < 0 # by default we load the entire dataset directly
test_labels = DATASET_INFO[baseset_name]['LABEL_GROUP']['pa_all'] - DATASET_INFO[baseset_name]['LABEL_GROUP'][_config["label_sets"]]
_log.info(f'###### Labels excluded in training : {[lb for lb in _config["exclude_cls_list"]]} ######')
_log.info(f'###### Unseen labels evaluated in testing: {[lb for lb in test_labels]} ######')
tr_parent = SuperpixelDataset( # base dataset
which_dataset = baseset_name,
base_dir=_config['path'][data_name]['data_dir'],
idx_split = _config['eval_fold'],
mode='train',
min_fg=str(_config["min_fg_data"]), # dummy entry for superpixel dataset
transforms=tr_transforms,
nsup = _config['task']['n_shots'],
scan_per_load = _config['scan_per_load'],
exclude_list = _config["exclude_cls_list"],
superpix_scale = _config["superpix_scale"],
fix_length = _config["max_iters_per_load"] if (data_name == 'C0_Superpix') or (data_name == 'CHAOST2_Superpix') else None
)
### dataloaders
trainloader = DataLoader(
tr_parent,
batch_size=_config['batch_size'],
shuffle=True,
num_workers=_config['num_workers'],
pin_memory=True,
drop_last=True
)
_log.info('###### Set optimizer ######')
if _config['optim_type'] == 'sgd':
optimizer = torch.optim.SGD(model.parameters(), **_config['optim'])
else:
raise NotImplementedError
scheduler = MultiStepLR(optimizer, milestones=_config['lr_milestones'], gamma = _config['lr_step_gamma'])
my_weight = compose_wt_simple(_config["use_wce"], data_name)
criterion = nn.CrossEntropyLoss(ignore_index=_config['ignore_label'], weight = my_weight)
i_iter = 0 # total number of iteration
n_sub_epoches = _config['n_steps'] // _config['max_iters_per_load'] # number of times for reloading
log_loss = {'loss': 0, 'align_loss': 0}
_log.info('###### Training ######')
for sub_epoch in range(n_sub_epoches):
_log.info(f'###### This is epoch {sub_epoch} of {n_sub_epoches} epoches ######')
for _, sample_batched in enumerate(trainloader):
# Prepare input
i_iter += 1
# add writers
support_images = [[shot.cuda() for shot in way]
for way in sample_batched['support_images']]
support_fg_mask = [[shot[f'fg_mask'].float().cuda() for shot in way]
for way in sample_batched['support_mask']]
support_bg_mask = [[shot[f'bg_mask'].float().cuda() for shot in way]
for way in sample_batched['support_mask']]
query_images = [query_image.cuda()
for query_image in sample_batched['query_images']]
query_labels = torch.cat(
[query_label.long().cuda() for query_label in sample_batched['query_labels']], dim=0)
optimizer.zero_grad()
# FIXME: in the model definition, filter out the failure case where pseudolabel falls outside of image or too small to calculate a prototype
try:
query_pred, align_loss, debug_vis, assign_mats = model(support_images, support_fg_mask, support_bg_mask, query_images, isval = False, val_wsize = None)
except:
print('Faulty batch detected, skip')
continue
query_loss = criterion(query_pred, query_labels)
loss = query_loss + align_loss
loss.backward()
optimizer.step()
scheduler.step()
# Log loss
query_loss = query_loss.detach().data.cpu().numpy()
align_loss = align_loss.detach().data.cpu().numpy() if align_loss != 0 else 0
_run.log_scalar('loss', query_loss)
_run.log_scalar('align_loss', align_loss)
log_loss['loss'] += query_loss
log_loss['align_loss'] += align_loss
# print loss and take snapshots
if (i_iter + 1) % _config['print_interval'] == 0:
loss = log_loss['loss'] / _config['print_interval']
align_loss = log_loss['align_loss'] / _config['print_interval']
log_loss['loss'] = 0
log_loss['align_loss'] = 0
print(f'step {i_iter+1}: loss: {loss}, align_loss: {align_loss},')
if (i_iter + 1) % _config['save_snapshot_every'] == 0:
_log.info('###### Taking snapshot ######')
torch.save(model.state_dict(),
os.path.join(f'{_run.observers[0].dir}/snapshots', f'{i_iter + 1}.pth'))
if data_name == 'C0_Superpix' or data_name == 'CHAOST2_Superpix':
if (i_iter + 1) % _config['max_iters_per_load'] == 0:
_log.info('###### Reloading dataset ######')
trainloader.dataset.reload_buffer()
print(f'###### New dataset with {len(trainloader.dataset)} slices has been loaded ######')
if (i_iter - 2) > _config['n_steps']:
return 1 # finish up