-
Notifications
You must be signed in to change notification settings - Fork 9
/
train_itm.py
387 lines (325 loc) · 17.6 KB
/
train_itm.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
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
import argparse
import os
import logging
import torch
import json
import random
import itertools
import time
import collections
import numpy as np
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader, ConcatDataset
from horovod import torch as hvd
from GLOBAL_VARIABLES import N_EXAMPLES_TEACHER
from uniter_model.data import ImageLmdbGroup
from uniter_model.data.loader import PrefetchLoader
from uniter_model.model.itm import UniterForImageTextRetrieval
from transformers.tokenization_bert import BertTokenizer
from dvl.options import default_params, add_itm_params, add_logging_params, add_kd_params, parse_with_config, map_db_dirs
from dvl.data.itm import TxtTokLmdb, ItmFastDataset, ItmValDataset, itm_fast_collate, itm_fast_collate_kd
from dvl.models.bi_encoder import BiEncoder, get_optimizer, setup_for_distributed_mode, \
BiEncoderNllLoss, get_schedule_linear, load_biencoder_checkpoint
from dvl.utils import print_args, num_of_parameters, _calc_loss, is_main_process
from dvl.hn import random_hard_neg, get_img_txt_mappings, sampled_hard_negatives
from dvl.const import IMG_DIM
from dvl.trainer import build_dataloader, _save_checkpoint, eval_model_on_dataloader, load_dataset
from dvl.indexer.faiss_indexers import DenseFlatIndexer
DEBUG_FLAG = False
torch.set_num_threads(4)
logger = logging.getLogger()
logger.setLevel(logging.INFO)
if logger.hasHandlers():
logger.handlers.clear()
console = logging.StreamHandler()
logger.addHandler(console)
def train_parser(parser):
default_params(parser)
add_itm_params(parser)
add_logging_params(parser)
add_kd_params(parser)
return parser
parser = argparse.ArgumentParser()
parser = train_parser(parser)
if DEBUG_FLAG:
args = parse_with_config(parser, [
'--config', './config/coco_ft_config_bert_debug.json',
'--sample_init_hard_negatives'
])
args.retrieval_mode = 'both'
else:
args = parse_with_config(parser)
if not os.path.exists(args.output_dir):
os.makedirs(args.output_dir, exist_ok=True)
# options safe guard
if args.conf_th == -1:
assert args.max_bb + args.max_txt_len + 2 <= 512
else:
assert args.num_bb + args.max_txt_len + 2 <= 512
hvd.init()
torch.cuda.set_device(hvd.local_rank())
args.device = torch.device("cuda", hvd.local_rank())
args.local_rank = hvd.rank()
args.n_gpu = hvd.size()
args.fp16_opt_level = 'O2' # for now let us assume always set opt level as O2
args.tokenizer = BertTokenizer.from_pretrained('bert-base-cased')
if args.project_dim > 0:
args.vector_size = args.project_dim
else:
args.vector_size = 768
if args.teacher_checkpoint is not None:
logger.info('teacher checkpoint is provided, using KD framework')
teacher_model = UniterForImageTextRetrieval.from_pretrained(os.path.join(args.teacher_checkpoint, 'config.json'),
state_dict=torch.load(os.path.join(args.teacher_checkpoint, 'model.pt')),
img_dim=IMG_DIM,
margin=0.2)
collate_func = itm_fast_collate_kd
else:
teacher_model = None
collate_func = itm_fast_collate
if is_main_process():
experiment = None
# options for DEBUG
print_args(args)
if experiment is not None:
experiment.log_parameters({
'train_batch_size': args.train_batch_size,
'learning_rate': args.learning_rate,
'num_hard_negatives': args.num_hard_negatives,
'hard_negatives_sampling': args.hard_negatives_sampling,
'caption_score_weight': args.caption_score_weight,
'kd_score_weight': args.kd_loss_weight,
'Temperature': args.T,
'project_dim': args.project_dim,
'retrieval_mode': args.retrieval_mode
})
if args.itm_global_file is not None:
with open(args.itm_global_file) as f:
args.img_meta = json.load(f)
else:
args.img_meta = None
# Init Model
bi_encoder = BiEncoder(args, args.fix_img_encoder, args.fix_txt_encoder, args.project_dim)
load_biencoder_checkpoint(bi_encoder, args.biencoder_checkpoint)
optimizer = get_optimizer(bi_encoder, args.learning_rate)
logger.info(f'total #params in img model = {num_of_parameters(bi_encoder.img_model)}, '
f'in txt model = {num_of_parameters(bi_encoder.txt_model)}')
logger.info(f'total #params in biencoder model = {num_of_parameters(bi_encoder, requires_grad=True)}')
bi_encoder, optimizer = setup_for_distributed_mode(bi_encoder, optimizer, args.device, args.n_gpu,
# args.local_rank,
-1,
args.fp16,
args.fp16_opt_level,
teacher_model=teacher_model)
# Load Data
all_img_dbs = ImageLmdbGroup(args.conf_th, args.max_bb, args.min_bb, args.num_bb, args.compressed_db)
# img2txt and txt2img mapping
train_img2txt, train_txt2img, train_img2set, train_txt2set, train_set2img, train_set2txt = \
get_img_txt_mappings(args.train_txt_dbs)
if args.sample_init_hard_negatives:
assert args.num_hard_negatives > 0, 'for init, num hard negatives has to > 0'
hard_neg_txt, hard_neg_img = sampled_hard_negatives(all_img_dbs, args, collate_func, bi_encoder, train_img2txt, train_txt2img)
else:
# hard_negatives = random_hard_neg(train_img2txt, args.num_hard_negatives, train_img2set, train_set2img)
if args.num_hard_negatives > 0:
raise NotImplementedError('random init hard negatives not impelmented yet')
else:
hard_neg_txt, hard_neg_img = None, None
val_img2txt = json.load(open(os.path.join(args.val_txt_db, 'img2txts.json')))
# load train and dev
logger.info(f"Loading Train Dataset "
f"{args.train_txt_dbs}, {args.train_img_dbs}")
train_dataset = load_dataset(all_img_dbs, args.train_txt_dbs, args.train_img_dbs, args, True)
for dset in train_dataset.datasets:
dset.new_epoch(hard_neg_img, hard_neg_txt)
train_dataloader = build_dataloader(train_dataset, collate_func, True, args)
logger.info(f'train dataset len = {len(train_dataset)}, dataloader len = {len(train_dataloader)}')
val_dataset = load_dataset(all_img_dbs, args.val_txt_db, args.val_img_db, args, is_train=False)
val_dataset.new_epoch()
val_dataloader = build_dataloader(val_dataset, collate_func, False, args)
logger.info(f'dev dataset len = {len(val_dataset)}, dataloader len = {len(val_dataloader)}')
updates_per_epoch = len(train_dataloader) // args.gradient_accumulation_steps
total_updates = updates_per_epoch * args.num_train_epochs
warmup_steps = int(0.1 * total_updates)
scheduler = get_schedule_linear(optimizer, warmup_steps, total_updates)
best_eval_metric = 0.0
if teacher_model:
# teacher model will always be in eval mode
teacher_model.eval()
for epoch in range(args.num_train_epochs):
epoch_loss, epoch_correct_predictions, rolling_train_loss = 0, 0, 0.0
bi_encoder.train()
logger.info('*' * 70)
if experiment is not None:
experiment.log_metric('epoch', epoch)
for dset in train_dataset.datasets:
dset.new_epoch(hard_neg_img, hard_neg_txt)
for step, batch in enumerate(train_dataloader):
model_out = bi_encoder(batch)
txt_vector, img_vectors, caption_vectors = model_out
loss_function = BiEncoderNllLoss()
bs = batch['sample_size']
if args.num_hard_negatives > 0:
loss_nce_txt, is_correct_txt, scores_txt = _calc_loss(args, loss_function, img_vectors[:bs], txt_vector, caption_vectors,
batch['pos_ctx_indices'], batch['neg_ctx_indices'], experiment)
loss_nce_img, is_correct_img, scores_img = _calc_loss(args, loss_function, txt_vector[:bs], img_vectors, caption_vectors,
batch['pos_ctx_indices'], batch['neg_ctx_indices'], experiment)
else:
loss_nce_txt, is_correct_txt, scores_txt = _calc_loss(args, loss_function, img_vectors, txt_vector,
caption_vectors,
batch['pos_ctx_indices'], batch['neg_ctx_indices'],
experiment)
loss_nce_img, is_correct_img, scores_img = _calc_loss(args, loss_function, txt_vector, img_vectors,
caption_vectors,
batch['pos_ctx_indices'], batch['neg_ctx_indices'],
experiment)
if args.retrieval_mode in['txt_only']:
raise ValueError('not supported anymore')
is_correct, scores = is_correct_txt.sum().item(), scores_txt
loss_nce = loss_nce_txt
elif args.retrieval_mode in ['img_only']:
raise ValueError('not supported anymore')
is_correct, scores = is_correct_img.sum().item(), scores_img
loss_nce = loss_nce_img
else:
is_correct = (is_correct_txt.sum().item() + is_correct_img.sum().item()) / 2
loss_nce = 0.5 * loss_nce_txt + 0.5 * loss_nce_img
scores = scores_txt * 0.5 + scores_img * 0.5
if teacher_model:
batch_new = {'gather_index': None}
for k in batch:
if 'teacher' in k:
new_k = k.replace('_teacher', '')
batch_new[new_k] = batch[k]
with torch.no_grad():
teacher_scores = teacher_model(batch_new, compute_loss=False).reshape(len(batch['txt_ids']), -1).T
assert teacher_scores.shape[0] == N_EXAMPLES_TEACHER, 'number of teacher example does not match'
# teacher_prob = teacher_prob.softmax(dim=1)
# KD loss
loss_kd = nn.KLDivLoss()(F.log_softmax(scores[:N_EXAMPLES_TEACHER] / args.T, dim=1),
F.softmax(teacher_scores / args.T, dim=1)) * args.T * args.T
loss = loss_nce + args.kd_loss_weight * loss_kd
else:
loss = loss_nce
if args.n_gpu > 1:
loss = loss.mean()
if args.gradient_accumulation_steps > 1:
loss = loss / args.gradient_accumulation_steps
epoch_correct_predictions += is_correct
epoch_loss += loss.item()
rolling_train_loss += loss.item()
if args.fp16:
from apex import amp
with amp.scale_loss(loss, optimizer) as scaled_loss:
scaled_loss.backward()
if args.max_grad_norm > 0:
torch.nn.utils.clip_grad_norm_(amp.master_params(optimizer), args.max_grad_norm)
else:
loss.backward()
if args.max_grad_norm > 0:
torch.nn.utils.clip_grad_norm_(bi_encoder.parameters(), args.max_grad_norm)
if (step+1) % args.log_result_step == 0 and is_main_process():
lr = optimizer.param_groups[0]['lr']
if args.teacher_checkpoint:
logger.info(
'Epoch: %d: Step: %d/%d, loss=%f, loss_nce=%f, loss_kd=%f, lr=%f', epoch, step,
len(train_dataloader), loss.item(), loss_nce.item(), loss_kd.item(), lr)
else:
logger.info(
'Epoch: %d: Step: %d/%d, loss=%f, loss_nce=%f, loss_kd=0.0, lr=%f', epoch, step,
len(train_dataloader), loss.item(), loss_nce.item(), lr)
if experiment is not None:
experiment.log_metric('step', step)
experiment.log_metric('lr', lr)
experiment.log_metric('loss_train', loss.item())
experiment.log_metric('loss_nce', loss_nce.item())
experiment.log_metric('loss_nce_txt', loss_nce_txt.item())
experiment.log_metric('loss_nce_img', loss_nce_img.item())
if args.teacher_checkpoint:
experiment.log_metric('loss_kd', loss_kd.item())
if (step + 1) % args.gradient_accumulation_steps == 0:
optimizer.step()
scheduler.step()
bi_encoder.zero_grad()
rolling_loss_step = args.log_result_step
if (step + 1) % rolling_loss_step == 0:
latest_rolling_train_av_loss = rolling_train_loss / rolling_loss_step
# logger.info('Train batch %d', step)
# logger.info('Avg. loss per last %d batches: %f', rolling_loss_step, latest_rolling_train_av_loss)
rolling_train_loss = 0.0
if experiment is not None and is_main_process():
experiment.log_metric('rolling_loss', latest_rolling_train_av_loss)
if step == 5 and DEBUG_FLAG:
logger.info('break for debug')
break
epoch_loss = (epoch_loss / len(train_dataloader)) if len(train_dataloader) > 0 else 0
total_samples = max(len(train_dataloader) * args.train_batch_size * 1, 1)
correct_ratio = float(epoch_correct_predictions / total_samples)
if experiment is not None and is_main_process():
logger.info(f'Av Loss per epoch = {epoch_loss}, epoch total correct predictions = {correct_ratio}')
experiment.log_metric('epoch_loss', epoch_loss)
experiment.log_metric('correct_number_train', epoch_correct_predictions)
experiment.log_metric('correct_ratio_train', correct_ratio)
# eval and save
bi_encoder.eval()
img2txt = dict(collections.ChainMap(*[json.load(open(os.path.join(db_folder, 'img2txts.json'))) for db_folder in [args.val_txt_db]]))
loss_val, correct_ratio_val, indexer_val, recall_both, _ = eval_model_on_dataloader(bi_encoder, val_dataloader, args, img2txt=img2txt)
recall_val = dict()
for t in recall_both[0]:
recall_val[t] = (recall_both[0][t] + recall_both[1][t]) / 2
current_eval_metric = np.mean(list(recall_val.values()))
if experiment is not None and is_main_process():
logger.info(f'val loss = {loss_val}. val correct prediction ratio = {correct_ratio_val}, recall = {recall_val}')
experiment.log_metric('total_valid_loss', loss_val)
experiment.log_metric('correct_ratio_valid', correct_ratio_val)
experiment.log_metric('R@1', recall_val[1])
experiment.log_metric('R@5', recall_val[5])
experiment.log_metric('R@10', recall_val[10])
experiment.log_metric('R@mean', current_eval_metric)
experiment.log_metric('img_R@1', recall_both[0][1])
experiment.log_metric('img_R@5', recall_both[0][5])
experiment.log_metric('img_R@10', recall_both[0][10])
experiment.log_metric('img_R@mean', np.mean(list(recall_both[0].values())))
experiment.log_metric('txt_R@1', recall_both[1][1])
experiment.log_metric('txt_R@5', recall_both[1][5])
experiment.log_metric('txt_R@10', recall_both[1][10])
experiment.log_metric('txt_R@mean', np.mean(list(recall_both[1].values())))
if current_eval_metric > best_eval_metric and is_main_process():
_save_checkpoint(args, bi_encoder, optimizer, scheduler, epoch, 0, 'best')
if is_main_process():
_save_checkpoint(args, bi_encoder, optimizer, scheduler, epoch, 0, 'last')
if is_main_process() and args.save_all_epochs:
_save_checkpoint(args, bi_encoder, optimizer, scheduler, epoch, 0)
# sample hard negative in here
if args.num_hard_negatives > 0:
hard_neg_txt, hard_neg_img = sampled_hard_negatives(all_img_dbs, args, collate_func, bi_encoder, train_img2txt,
train_txt2img)
else:
# no hard negative sampling
hard_neg_txt, hard_neg_img = None, None
assert args.hard_negatives_sampling == 'none', f'sampleing method {args.hard_negatives_sampling} is not none'
if args.test_txt_db:
test_dataset = load_dataset(all_img_dbs, args.test_txt_db, args.test_img_db, args, is_train=False)
test_dataset.new_epoch()
test_dataloader = build_dataloader(test_dataset, collate_func, False, args)
logger.info(f'test dataset len = {len(test_dataset)}, dataloader len = {len(test_dataloader)}')
bi_encoder.eval()
img2txt = dict(collections.ChainMap(*[json.load(open(os.path.join(db_folder, 'img2txts.json'))) for db_folder in [args.test_txt_db]]))
loss_test, correct_ratio_test, indexer_test, (recall_img, recall_txt), _ = eval_model_on_dataloader(bi_encoder, test_dataloader, args,
args.txt_retrieval, img2txt)
recall_mean = np.mean(list(recall_img.values()))
if experiment is not None:
experiment.log_metric('test_img:R@1', "{:.4f}".format(round(recall_img[1], 4)))
experiment.log_metric('test_img:R@5', "{:.4f}".format(round(recall_img[5], 4)))
experiment.log_metric('test_img:R@10', "{:.4f}".format(round(recall_img[10], 4)))
experiment.log_metric('test_img:R@mean', "{:.4f}".format(round(recall_mean, 4)))
recall_mean = np.mean(list(recall_txt.values()))
if experiment is not None:
experiment.log_metric('test_txt:R@1', "{:.4f}".format(round(recall_txt[1], 4)))
experiment.log_metric('test_txt:R@5', "{:.4f}".format(round(recall_txt[5], 4)))
experiment.log_metric('test_txt:R@10', "{:.4f}".format(round(recall_txt[10], 4)))
experiment.log_metric('test_txt:R@mean', "{:.4f}".format(round(recall_mean, 4)))
experiment.log_metric('correct_ratio_test', "{:.4f}".format(round(correct_ratio_test, 4)))
experiment.log_metric('loss_test', "{:.4f}".format(loss_test))
experiment.log_metric('n_image_test', len(indexer_test.index_id_to_db_id))