def main(opts):
hvd.init()
if hvd.rank() == 0:
toker = RobertaTokenizer.from_pretrained('roberta-base')
all_gather_list(None)
else:
all_gather_list(None)
toker = RobertaTokenizer.from_pretrained('roberta-base')
model_opts = Struct(json.load(open(f"{opts.model_dir}/log/hps.json")))
model_config = f"{opts.model_dir}/log/model_config.json"
video_db = load_video_sub_dataset(model_opts.vfeat_db,
model_opts.sub_txt_db,
model_opts.vfeat_interval,
model_opts)
dset = TvcEvalDataset(video_db, opts.target_clip)
loader = build_dataloader(dset, opts.batch_size,
TvcEvalDataset.collate, False, opts)
checkpoint = torch.load(f"{opts.model_dir}/ckpt/"
f"model_step_{opts.ckpt_step}.pt")
img_pos_embed_weight_key = "v_encoder.f_encoder.img_embeddings" +\
".position_embeddings.weight"
if img_pos_embed_weight_key in checkpoint:
max_frm_seq_len = len(checkpoint[img_pos_embed_weight_key])
else:
max_frm_seq_len = MAX_FRM_SEQ_LEN
model = HeroForTvc.from_pretrained(model_config,
state_dict=checkpoint,
vfeat_dim=VFEAT_DIM,
max_frm_seq_len=max_frm_seq_len,
lsr=model_opts.lsr)
model.cuda()
model = amp.initialize(model, enabled=opts.fp16, opt_level='O2')
bos = toker.convert_tokens_to_ids(['<s>'])[0]
eos = toker.convert_tokens_to_ids(['</s>'])[0]
model.eval()
generator = TvcGenerator(model, opts.max_gen_step, bos, eos, opts.fp16)
results = decode(loader, generator, toker)
save_jsonl(results, opts.output)
# evaluate score if possible
if (hvd.rank() == 0
and 'descs' in json.loads(next(iter(open(opts.target_clip))))):
evaluator = TVCEval(opts.target_clip)
score = evaluator(results)
print(score)
def build_dataloader(opts):
# Load ground truth, query db and video db
hps_file = f'{opts.output_dir}/log/hps.json'
model_opts = Struct(load_json(hps_file))
video_ids = get_video_ids(opts.query_txt_db)
video_db = load_video_sub_dataset(opts.vfeat_db, opts.sub_txt_db,
model_opts.vfeat_interval, model_opts)
assert opts.split in opts.query_txt_db, (opts.split, opts.query_txt_db)
q_txt_db = QueryTokLmdb(opts.query_txt_db, -1)
eval_dataset = VcmrFullEvalDataset(video_ids,
video_db,
q_txt_db,
distributed=model_opts.distributed_eval)
eval_dataloader = DataLoader(eval_dataset,
batch_size=opts.batch_size,
num_workers=opts.n_workers,
pin_memory=opts.pin_mem,
collate_fn=vcmr_full_eval_collate)
eval_dataloader = PrefetchLoader(eval_dataloader)
return eval_dataloader
def main(opts):
hvd.init()
n_gpu = hvd.size()
device = torch.device("cuda", hvd.local_rank())
torch.cuda.set_device(hvd.local_rank())
rank = hvd.rank()
opts.rank = rank
LOGGER.info("device: {} n_gpu: {}, rank: {}, "
"16-bits training: {}".format(device, n_gpu, hvd.rank(),
opts.fp16))
if hvd.rank() != 0:
LOGGER.disabled = True
if opts.gradient_accumulation_steps < 1:
raise ValueError("Invalid gradient_accumulation_steps parameter: {}, "
"should be >= 1".format(
opts.gradient_accumulation_steps))
set_random_seed(opts.seed)
opts.task = 'tvc'
# train_examples = None
LOGGER.info(f"Loading the whole video dataset {opts.sub_txt_db}, "
f"{opts.vfeat_db}")
video_db = load_video_sub_dataset(opts.vfeat_db, opts.sub_txt_db,
opts.vfeat_interval, opts)
# data loaders
# train
LOGGER.info(f"Loading train dataset {opts.train_db}")
train_cap = CaptionTokLmdb(opts.train_db, opts.max_txt_len)
train_dset = TvcTrainDataset(video_db, train_cap, opts.max_cap_per_vid)
LOGGER.info(f"{sum(all_gather_list(len(train_dset)))} samples loaded")
train_loader = build_dataloader(train_dset, opts.train_batch_size,
TvcTrainDataset.collate, True, opts)
# val
LOGGER.info(f"Loading val dataset {opts.val_db}")
val_cap = CaptionTokLmdb(opts.val_db, -1)
val_dset = TvcValDataset(video_db, val_cap, -1)
val_loader = build_dataloader(val_dset, opts.val_batch_size,
TvcValDataset.collate, False, opts)
if hvd.rank() == 0:
evaluator = TVCEval(opts.val_ref)
else:
evaluator = NoOp()
# Prepare model
if opts.checkpoint:
checkpoint = torch.load(opts.checkpoint)
else:
checkpoint = {}
img_pos_embed_weight_key = "v_encoder.f_encoder.img_embeddings" +\
".position_embeddings.weight"
if img_pos_embed_weight_key in checkpoint:
max_frm_seq_len = len(checkpoint[img_pos_embed_weight_key])
else:
max_frm_seq_len = MAX_FRM_SEQ_LEN
model = HeroForTvc.from_pretrained(opts.model_config,
state_dict=checkpoint,
vfeat_dim=VFEAT_DIM,
max_frm_seq_len=max_frm_seq_len,
lsr=opts.lsr)
model.to(device)
# make sure every process has same model parameters in the beginning
broadcast_tensors([p.data for p in model.parameters()], 0)
set_dropout(model, opts.dropout)
# Prepare optimizer
optimizer = build_optimizer(model, opts)
model, optimizer = amp.initialize(model,
optimizer,
enabled=opts.fp16,
opt_level='O2')
# assumes roberta tokenizer only
if hvd.local_rank() == 0:
# quick hack to prevent multi-process download collision
toker = RobertaTokenizer.from_pretrained('roberta-base')
all_gather_list(None)
else:
all_gather_list(None)
toker = RobertaTokenizer.from_pretrained('roberta-base')
bos = toker.convert_tokens_to_ids(['<s>'])[0]
eos = toker.convert_tokens_to_ids(['</s>'])[0]
generator = TvcGenerator(model, opts.max_gen_step, bos, eos, opts.fp16)
global_step = 0
if rank == 0:
save_training_meta(opts)
TB_LOGGER.create(join(opts.output_dir, 'log'))
pbar = tqdm(total=opts.num_train_steps)
model_saver = ModelSaver(join(opts.output_dir, 'ckpt'))
os.makedirs(join(opts.output_dir, 'results')) # store val predictions
add_log_to_file(join(opts.output_dir, 'log', 'log.txt'))
else:
LOGGER.disabled = True
pbar = NoOp()
model_saver = NoOp()
LOGGER.info(f"***** Running training with {n_gpu} GPUs *****")
LOGGER.info(" Batch size = %d", opts.train_batch_size)
LOGGER.info(" Accumulate steps = %d", opts.gradient_accumulation_steps)
LOGGER.info(" Num steps = %d", opts.num_train_steps)
train_loss = RunningMeter('loss')
n_vid = 0
n_cap = 0
n_epoch = 0
start = time()
# quick hack for amp delay_unscale bug
optimizer.zero_grad()
optimizer.step()
model.train()
while True:
for step, batch in enumerate(train_loader):
n_vid += opts.train_batch_size
n_cap += batch['cap_input_ids'].size(0)
loss = model(batch, compute_loss=True)
loss = loss.mean()
train_loss(loss.item())
delay_unscale = (step + 1) % opts.gradient_accumulation_steps != 0
with amp.scale_loss(loss, optimizer,
delay_unscale=delay_unscale) as scaled_loss:
scaled_loss.backward()
if not delay_unscale:
# gather gradients from every processes
# do this before unscaling to make sure every process uses
# the same gradient scale
grads = [
p.grad.data for p in model.parameters()
if p.requires_grad and p.grad is not None
]
all_reduce_and_rescale_tensors(grads, float(1))
if (step + 1) % opts.gradient_accumulation_steps == 0:
global_step += 1
# learning rate scheduling
lr_this_step = get_lr_sched(global_step, opts)
for i, param_group in enumerate(optimizer.param_groups):
if i == 0 or i == 1:
param_group['lr'] = lr_this_step * opts.lr_mul
elif i == 2 or i == 3:
param_group['lr'] = lr_this_step
else:
raise ValueError()
TB_LOGGER.add_scalar('lr', lr_this_step, global_step)
# log loss
TB_LOGGER.add_scalar(train_loss.name, train_loss.val,
global_step)
TB_LOGGER.step()
# update model params
if opts.grad_norm != -1:
grad_norm = clip_grad_norm_(amp.master_params(optimizer),
opts.grad_norm)
TB_LOGGER.add_scalar('grad_norm', grad_norm, global_step)
optimizer.step()
optimizer.zero_grad()
pbar.update(1)
if global_step % 100 == 0:
# monitor training throughput
LOGGER.info('-------------------------------------------')
LOGGER.info(f'Step {global_step}:')
tot_vid = sum(all_gather_list(n_vid))
vid_per_sec = int(tot_vid / (time() - start))
LOGGER.info(f'{tot_vid} videos trained at '
f'{vid_per_sec} vid/s')
tot_cap = sum(all_gather_list(n_cap))
cap_per_sec = int(tot_cap / (time() - start))
TB_LOGGER.add_scalar(f'perf/vid_per_s', vid_per_sec,
global_step)
TB_LOGGER.add_scalar(f'perf/cap_per_s', cap_per_sec,
global_step)
if global_step % opts.valid_steps == 0:
LOGGER.info('===========================================')
LOGGER.info(f"Step {global_step}: start validation")
val_log, results = validate(val_loader, generator, toker,
evaluator)
if hvd.rank() == 0:
save_jsonl(
results, f"{opts.output_dir}/results/"
f"/results_{global_step}.jsonl")
TB_LOGGER.log_scaler_dict(val_log)
LOGGER.info('===========================================')
model_saver.save(model, global_step)
if global_step >= opts.num_train_steps:
break
n_epoch += 1
LOGGER.info(f"finished {n_epoch} epochs")
if global_step >= opts.num_train_steps:
break
LOGGER.info('===========================================')
if global_step % opts.valid_steps != 0:
val_log, results = validate(val_loader, generator, toker, evaluator)
if hvd.rank() == 0:
save_jsonl(
results, f"{opts.output_dir}/results/"
f"/results_{global_step}.jsonl")
TB_LOGGER.log_scaler_dict(val_log)
model_saver.save(model, global_step)
def main(opts):
hvd.init()
n_gpu = hvd.size()
device = torch.device("cuda", hvd.local_rank())
torch.cuda.set_device(hvd.local_rank())
opts.n_gpu = n_gpu
LOGGER.info("device: {} n_gpu: {}, rank: {}, "
"16-bits training: {}".format(device, n_gpu, hvd.rank(),
opts.fp16))
if hvd.rank() != 0:
LOGGER.disabled = True
set_random_seed(opts.seed)
# train_examples = None
LOGGER.info(f"Loading the whole video dataset {opts.sub_txt_db}, "
f"{opts.vfeat_db}")
if opts.task != "didemo_video_only":
video_db = load_video_sub_dataset(opts.vfeat_db, opts.sub_txt_db,
opts.vfeat_interval, opts)
else:
txt_meta = load_json(join(opts.train_query_txt_db, "meta.json"))
video_db = load_video_only_dataset(opts.vfeat_db, txt_meta,
opts.vfeat_interval, opts)
# data loaders
# train
video_ids = get_video_ids(opts.train_query_txt_db)
train_q_txt_db = QueryTokLmdb(opts.train_query_txt_db, opts.max_txt_len)
train_dataloaders = build_downstream_dataloaders([opts.task],
video_db,
video_ids,
True,
opts,
shuffle=True,
q_txt_db=train_q_txt_db)
meta_loader = MetaLoader(train_dataloaders,
accum_steps=opts.gradient_accumulation_steps,
distributed=n_gpu > 1)
meta_loader = PrefetchLoader(meta_loader)
# val
video_ids = get_video_ids(opts.val_query_txt_db)
val_q_txt_db = QueryTokLmdb(opts.val_query_txt_db, -1)
val_dataloaders = build_downstream_dataloaders([opts.task],
video_db,
video_ids,
False,
opts,
q_txt_db=val_q_txt_db)
if opts.task != "didemo_video_only":
inf_dataset = VcmrFullEvalDataset
else:
inf_dataset = VcmrVideoOnlyFullEvalDataset
LOGGER.info(f"Loading Inference Dataset {opts.val_query_txt_db} (val)")
val_dset = inf_dataset(video_ids,
video_db,
val_q_txt_db,
distributed=opts.distributed_eval)
inf_loader_val = DataLoader(val_dset,
batch_size=opts.vcmr_eval_q_batch_size,
num_workers=opts.n_workers,
pin_memory=opts.pin_mem,
collate_fn=vcmr_full_eval_collate)
inf_loader_val = PrefetchLoader(inf_loader_val)
if opts.test_query_txt_db:
LOGGER.info(
f"Loading Inference Dataset {opts.test_query_txt_db} (test)")
video_ids = get_video_ids(opts.test_query_txt_db)
test_q_txt_db = QueryTokLmdb(opts.test_query_txt_db, -1)
test_dset = inf_dataset(video_ids,
video_db,
test_q_txt_db,
distributed=opts.distributed_eval)
inf_loader_test = DataLoader(test_dset,
batch_size=opts.vcmr_eval_q_batch_size,
num_workers=opts.n_workers,
pin_memory=opts.pin_mem,
collate_fn=vcmr_full_eval_collate)
inf_loader_test = PrefetchLoader(inf_loader_test)
# Prepare model
if opts.checkpoint:
checkpoint = torch.load(opts.checkpoint)
else:
checkpoint = {}
img_pos_embed_weight_key = "v_encoder.f_encoder.img_embeddings" +\
".position_embeddings.weight"
if img_pos_embed_weight_key in checkpoint:
max_frm_seq_len = len(checkpoint[img_pos_embed_weight_key])
else:
max_frm_seq_len = MAX_FRM_SEQ_LEN
model = HeroForVcmr.from_pretrained(
opts.model_config,
state_dict=checkpoint,
vfeat_dim=VFEAT_DIM,
max_frm_seq_len=max_frm_seq_len,
lw_neg_ctx=opts.lw_neg_ctx,
lw_neg_q=opts.lw_neg_q,
lw_st_ed=0,
ranking_loss_type=opts.ranking_loss_type,
use_hard_negative=False,
hard_pool_size=opts.hard_pool_size,
margin=opts.margin,
use_all_neg=opts.use_all_neg,
drop_svmr_prob=opts.drop_svmr_prob)
model.to(device)
# make sure every process has same model parameters in the beginning
broadcast_tensors([p.data for p in model.parameters()], 0)
set_dropout(model, opts.dropout)
# Prepare optimizer
optimizer = build_optimizer(model, opts)
task2scaler = {t: i for i, t in enumerate(train_dataloaders.keys())}
model, optimizer = amp.initialize(model,
optimizer,
num_losses=len(task2scaler),
enabled=opts.fp16,
opt_level='O2')
restorer = TrainingRestorer(opts, model, optimizer)
global_step = restorer.global_step
TB_LOGGER.global_step = global_step
if hvd.rank() == 0:
save_training_meta(opts)
TB_LOGGER.create(join(opts.output_dir, 'log'))
pbar = tqdm(total=opts.num_train_steps)
model_saver = ModelSaver(join(opts.output_dir, 'ckpt'))
if not exists(join(opts.output_dir, 'results')):
# store tvr predictions
os.makedirs(join(opts.output_dir, 'results'))
if opts.nms_thd != -1:
# store tvr-nms predictions
if not exists(join(opts.output_dir, 'results_nms')):
os.makedirs(join(opts.output_dir, 'results_nms'))
add_log_to_file(join(opts.output_dir, 'log', 'log.txt'))
else:
pbar = NoOp()
model_saver = NoOp()
restorer = NoOp()
if global_step > 0:
pbar.update(global_step)
LOGGER.info(f"***** Running training with {n_gpu} GPUs *****")
LOGGER.info(" Batch size = %d", opts.train_batch_size)
LOGGER.info(" Accumulate steps = %d", opts.gradient_accumulation_steps)
LOGGER.info(" Num steps = %d", opts.num_train_steps)
task2loss = {
task: RunningMeter(f'loss/{task}')
for task in train_dataloaders.keys()
}
for obj in (f'{opts.task}_st_ed', f'{opts.task}_neg_ctx',
f'{opts.task}_neg_q'):
task2loss[obj] = RunningMeter(f'loss/{obj}')
model.train()
n_examples = defaultdict(int)
start = time()
# quick hack for amp delay_unscale bug
optimizer.zero_grad()
if global_step == 0:
optimizer.step()
for step, (task, batch) in enumerate(meta_loader):
if len(opts.hard_negtiave_start_step) > 0:
for i, hn_step in enumerate(opts.hard_negtiave_start_step):
if global_step >= hn_step and hn_step != -1:
model.set_hard_negative(True, opts.hard_pool_size[i],
opts.hard_neg_weights[i])
if opts.train_span_start_step != -1 and\
global_step >= opts.train_span_start_step:
model.set_train_st_ed(opts.lw_st_ed)
n_examples[task] += opts.train_batch_size
loss = model(batch, task=task, compute_loss=True)
loss_st_ed, loss_neg_ctx, loss_neg_q = loss
loss = loss_st_ed + loss_neg_ctx + loss_neg_q
for n, ls, w in (('st_ed', loss_st_ed, opts.lw_st_ed),
('neg_ctx', loss_neg_ctx, opts.lw_neg_ctx),
('neg_q', loss_neg_q, opts.lw_neg_q)):
ls = ls.item()
if w:
ls /= w
task2loss[f'{task}_{n}'](ls)
loss = loss.mean()
task2loss[task](loss.item())
delay_unscale = (step + 1) % opts.gradient_accumulation_steps != 0
with amp.scale_loss(loss,
optimizer,
delay_unscale=delay_unscale,
loss_id=task2scaler[task]) as scaled_loss:
scaled_loss.backward()
if not delay_unscale:
# gather gradients from every processes
# do this before unscaling to make sure every process uses
# the same gradient scale
grads = [
p.grad.data for p in model.parameters()
if p.requires_grad and p.grad is not None
]
all_reduce_and_rescale_tensors(grads, float(1))
if (step + 1) % opts.gradient_accumulation_steps == 0:
global_step += 1
# learning rate scheduling
lr_this_step = get_lr_sched(global_step, opts)
for param_group in optimizer.param_groups:
param_group['lr'] = lr_this_step
TB_LOGGER.add_scalar('lr', lr_this_step, global_step)
# log loss
TB_LOGGER.log_scaler_dict({
temp_loss.name: temp_loss.val
for temp_loss in task2loss.values()
if temp_loss.val is not None
})
TB_LOGGER.step()
# update model params
if opts.grad_norm != -1:
grad_norm = clip_grad_norm_(amp.master_params(optimizer),
opts.grad_norm)
TB_LOGGER.add_scalar('grad_norm', grad_norm, global_step)
optimizer.step()
optimizer.zero_grad()
pbar.update(1)
if global_step % 100 == 0:
# monitor training throughput
LOGGER.info('-------------------------------------------')
LOGGER.info(f'Step {global_step}:')
for t in train_dataloaders.keys():
tot_ex = sum(all_gather_list(n_examples[t]))
ex_per_sec = int(tot_ex / (time() - start))
LOGGER.info(f'{t}: {tot_ex} examples trained at '
f'{ex_per_sec} ex/s')
TB_LOGGER.add_scalar(f'perf/{t}_ex_per_s', ex_per_sec,
global_step)
if global_step % opts.valid_steps == 0:
LOGGER.info('===========================================')
LOGGER.info(f"Step {global_step}: start running validation")
validate(model, val_dataloaders, opts)
if hvd.rank() == 0 or opts.distributed_eval:
log, results = validate_full_vcmr(model,
inf_loader_val,
'val',
opts,
model_opts=opts)
save_json(
results, f'{opts.output_dir}/results/'
f'val_results_{global_step}_rank{hvd.rank()}.json')
TB_LOGGER.log_scaler_dict(log)
if opts.test_query_txt_db:
log, results = validate_full_vcmr(model,
inf_loader_test,
'test',
opts,
model_opts=opts)
save_json(
results, f'{opts.output_dir}/results/'
f'test_results_{global_step}_rank{hvd.rank()}.json'
)
TB_LOGGER.log_scaler_dict(log)
LOGGER.info('===========================================')
model_saver.save(model, global_step)
# step restorer in the end to prevent missing validation checkpoint
restorer.step()
if global_step >= opts.num_train_steps:
break
LOGGER.info('===========================================')
if global_step % opts.valid_steps != 0:
if hvd.rank() == 0 or opts.distributed_eval:
log, results = validate_full_vcmr(model,
inf_loader_val,
'val',
opts,
model_opts=opts)
save_json(
results, f'{opts.output_dir}/results/'
f'val_results_{global_step}'
f'_rank{hvd.rank()}_final.json')
TB_LOGGER.log_scaler_dict(log)
if opts.test_query_txt_db:
log, results = validate_full_vcmr(model,
inf_loader_test,
'test',
opts,
model_opts=opts)
save_json(
results, f'{opts.output_dir}/results/'
f'test_results_{global_step}_rank{hvd.rank()}.json')
TB_LOGGER.log_scaler_dict(log)
model_saver.save(model, f'{global_step}_final')
def main(opts):
hvd.init()
n_gpu = hvd.size()
device = torch.device("cuda", hvd.local_rank())
torch.cuda.set_device(hvd.local_rank())
LOGGER.info("device: {} n_gpu: {}, rank: {}, "
"16-bits training: {}".format(device, n_gpu, hvd.rank(),
opts.fp16))
if hvd.rank() != 0:
LOGGER.disabled = True
hps_file = f'{opts.output_dir}/log/hps.json'
model_opts = Struct(json.load(open(hps_file)))
model_config = f'{opts.output_dir}/log/model_config.json'
# load DBs and image dirs
video_ids = get_video_ids(opts.query_txt_db)
video_db = load_video_sub_dataset(opts.vfeat_db, opts.sub_txt_db,
model_opts.vfeat_interval, model_opts)
assert opts.split in opts.query_txt_db
q_txt_db = QaQueryTokLmdb(opts.query_txt_db, -1)
eval_dataset = ViolinEvalDataset(video_ids,
video_db,
q_txt_db,
sampled_by_q=model_opts.sampled_by_q)
collate_fn = violin_eval_collate
# Prepare model
if exists(opts.checkpoint):
ckpt_file = opts.checkpoint
else:
ckpt_file = f'{opts.output_dir}/ckpt/model_step_{opts.checkpoint}.pt'
checkpoint = torch.load(ckpt_file)
img_pos_embed_weight_key = "v_encoder.f_encoder.img_embeddings" +\
".position_embeddings.weight"
assert img_pos_embed_weight_key in checkpoint
max_frm_seq_len = len(checkpoint[img_pos_embed_weight_key])
model = HeroForViolin.from_pretrained(model_config,
state_dict=checkpoint,
vfeat_dim=VFEAT_DIM,
max_frm_seq_len=max_frm_seq_len)
model.to(device)
if opts.fp16:
model = amp.initialize(model, enabled=opts.fp16, opt_level='O2')
eval_dataloader = DataLoader(eval_dataset,
batch_size=opts.batch_size,
num_workers=opts.n_workers,
pin_memory=opts.pin_mem,
collate_fn=collate_fn)
eval_dataloader = PrefetchLoader(eval_dataloader)
_, results, logits = validate_violin(model, eval_dataloader, opts.split,
opts.save_logits)
result_dir = f'{opts.output_dir}/results_{opts.split}'
if opts.save_logits:
result_dir += '_w_logit'
if not exists(result_dir) and hvd.rank() == 0:
os.makedirs(result_dir)
all_results = {}
for id2res in all_gather_list(results):
all_results.update(id2res)
if opts.save_logits:
all_logits = {}
for id2logit in all_gather_list(logits):
all_logits.update(id2logit)
if hvd.rank() == 0:
save_json(all_results,
f'{result_dir}/results_{opts.checkpoint}_all.json')
LOGGER.info('All results written......')
if opts.save_logits:
save_pickle(all_logits,
f'{result_dir}/logits_{opts.checkpoint}_all.pkl')
LOGGER.info('All logits written......')
def main(opts):
hvd.init()
n_gpu = hvd.size()
device = torch.device("cuda", hvd.local_rank())
torch.cuda.set_device(hvd.local_rank())
opts.n_gpu = n_gpu
LOGGER.info("device: {} n_gpu: {}, rank: {}, "
"16-bits training: {}".format(device, n_gpu, hvd.rank(),
opts.fp16))
if hvd.rank() != 0:
LOGGER.disabled = True
set_random_seed(opts.seed)
# train_examples = None
LOGGER.info(f"Loading the whole video dataset {opts.sub_txt_db}, "
f"{opts.vfeat_db}")
video_db = load_video_sub_dataset(opts.vfeat_db, opts.sub_txt_db,
opts.vfeat_interval, opts)
# data loaders
# train
LOGGER.info(f"Loading the train QA dataset {opts.train_query_txt_db}")
video_ids = get_video_ids(opts.train_query_txt_db)
train_q_txt_db = QaQueryTokLmdb(opts.train_query_txt_db, opts.max_txt_len)
train_dataloaders = build_downstream_dataloaders([opts.task],
video_db,
video_ids,
True,
opts,
q_txt_db=train_q_txt_db,
shuffle=True)
meta_loader = MetaLoader(train_dataloaders,
accum_steps=opts.gradient_accumulation_steps,
distributed=n_gpu > 1)
meta_loader = PrefetchLoader(meta_loader)
# val
LOGGER.info(f"Loading the val QA dataset {opts.val_query_txt_db}")
video_ids = get_video_ids(opts.val_query_txt_db)
val_q_txt_db = QaQueryTokLmdb(opts.val_query_txt_db, -1)
val_dataloaders = build_downstream_dataloaders([opts.task],
video_db,
video_ids,
False,
opts,
q_txt_db=val_q_txt_db)
if opts.test_query_txt_db:
LOGGER.info(f"Loading the test QA dataset {opts.test_query_txt_db}")
video_ids = get_video_ids(opts.test_query_txt_db)
test_q_txt_db = QaQueryTokLmdb(opts.test_query_txt_db, -1)
test_dataloaders = build_downstream_dataloaders([opts.task],
video_db,
video_ids,
False,
opts,
q_txt_db=test_q_txt_db)
# Prepare model
if opts.checkpoint:
checkpoint = torch.load(opts.checkpoint)
else:
checkpoint = {}
img_pos_embed_weight_key = "v_encoder.f_encoder.img_embeddings" +\
".position_embeddings.weight"
if img_pos_embed_weight_key in checkpoint:
max_frm_seq_len = len(checkpoint[img_pos_embed_weight_key])
else:
max_frm_seq_len = MAX_FRM_SEQ_LEN
model = HeroForVideoQA.from_pretrained(opts.model_config,
state_dict=checkpoint,
vfeat_dim=VFEAT_DIM,
max_frm_seq_len=max_frm_seq_len)
model.to(device)
# make sure every process has same model parameters in the beginning
broadcast_tensors([p.data for p in model.parameters()], 0)
set_dropout(model, opts.dropout)
# Prepare optimizer
optimizer = build_optimizer(model, opts)
task2scaler = {t: i for i, t in enumerate(train_dataloaders.keys())}
model, optimizer = amp.initialize(model,
optimizer,
num_losses=len(task2scaler),
enabled=opts.fp16,
opt_level='O2')
restorer = TrainingRestorer(opts, model, optimizer)
global_step = restorer.global_step
TB_LOGGER.global_step = global_step
if hvd.rank() == 0:
save_training_meta(opts)
TB_LOGGER.create(join(opts.output_dir, 'log'))
pbar = tqdm(total=opts.num_train_steps)
model_saver = ModelSaver(join(opts.output_dir, 'ckpt'))
if not exists(join(opts.output_dir, 'results')):
# store tvqa predictions
os.makedirs(join(opts.output_dir, 'results'))
add_log_to_file(join(opts.output_dir, 'log', 'log.txt'))
else:
LOGGER.disabled = True
pbar = NoOp()
model_saver = NoOp()
restorer = NoOp()
if global_step > 0:
pbar.update(global_step)
LOGGER.info(f"***** Running training with {n_gpu} GPUs *****")
LOGGER.info(" Batch size = %d", opts.train_batch_size)
LOGGER.info(" Accumulate steps = %d", opts.gradient_accumulation_steps)
LOGGER.info(" Num steps = %d", opts.num_train_steps)
task2loss = {
task: RunningMeter(f'loss/{task}')
for task in train_dataloaders.keys()
}
for obj in (f'{opts.task}_qa', f'{opts.task}_st_ed'):
task2loss[obj] = RunningMeter(f'loss/{obj}')
model.train()
n_examples = defaultdict(int)
start = time()
# quick hack for amp delay_unscale bug
optimizer.zero_grad()
if global_step == 0:
optimizer.step()
for step, (task, batch) in enumerate(meta_loader):
n_examples[task] += opts.train_batch_size
loss = model(batch, task=task, compute_loss=True)
loss_qa, loss_st_ed = loss
loss = loss_qa + opts.lw_st_ed * loss_st_ed
for n, ls in (('st_ed', loss_st_ed), ('qa', loss_qa)):
ls = ls.item()
task2loss[f'{task}_{n}'](ls)
loss = loss.mean()
task2loss[task](loss.item())
delay_unscale = (step + 1) % opts.gradient_accumulation_steps != 0
with amp.scale_loss(loss,
optimizer,
delay_unscale=delay_unscale,
loss_id=task2scaler[task]) as scaled_loss:
scaled_loss.backward()
if not delay_unscale:
# gather gradients from every processes
# do this before unscaling to make sure every process uses
# the same gradient scale
grads = [
p.grad.data for p in model.parameters()
if p.requires_grad and p.grad is not None
]
all_reduce_and_rescale_tensors(grads, float(1))
if (step + 1) % opts.gradient_accumulation_steps == 0:
global_step += 1
# learning rate scheduling
lr_this_step = get_lr_sched(global_step, opts)
for i, param_group in enumerate(optimizer.param_groups):
if i == 0 or i == 1:
param_group['lr'] = lr_this_step * opts.lr_mul
elif i == 2 or i == 3:
param_group['lr'] = lr_this_step
else:
raise ValueError()
TB_LOGGER.add_scalar('lr', lr_this_step, global_step)
TB_LOGGER.log_scaler_dict({
temp_loss.name: temp_loss.val
for temp_loss in task2loss.values()
if temp_loss.val is not None
})
TB_LOGGER.step()
# update model params
if opts.grad_norm != -1:
grad_norm = clip_grad_norm_(amp.master_params(optimizer),
opts.grad_norm)
TB_LOGGER.add_scalar('grad_norm', grad_norm, global_step)
optimizer.step()
optimizer.zero_grad()
restorer.step()
pbar.update(1)
if global_step % 100 == 0:
# monitor training throughput
LOGGER.info('-------------------------------------------')
LOGGER.info(f'Step {global_step}:')
for t in train_dataloaders.keys():
tot_ex = sum(all_gather_list(n_examples[t]))
ex_per_sec = int(tot_ex / (time() - start))
LOGGER.info(f'{t}: {tot_ex} examples trained at '
f'{ex_per_sec} ex/s')
TB_LOGGER.add_scalar(f'perf/{t}_ex_per_s', ex_per_sec,
global_step)
if global_step % opts.valid_steps == 0:
LOGGER.info('===========================================')
LOGGER.info(f"Step {global_step}: start running validation")
validate(model,
val_dataloaders,
"val",
opts,
global_step=global_step)
if opts.test_query_txt_db:
validate(model,
test_dataloaders,
"test",
opts,
global_step=global_step)
LOGGER.info('===========================================')
model_saver.save(model, global_step)
if global_step >= opts.num_train_steps:
break
LOGGER.info('===========================================')
if global_step % opts.valid_steps != 0:
LOGGER.info('===========================================')
LOGGER.info(f"Step {global_step}: start running validation")
validate(model, val_dataloaders, "val", opts, global_step=global_step)
if opts.test_query_txt_db:
validate(model,
test_dataloaders,
"test",
opts,
global_step=global_step)
LOGGER.info('===========================================')
model_saver.save(model, f'{global_step}_final')
def main(opts):
hvd.init()
n_gpu = hvd.size()
device = torch.device("cuda", hvd.local_rank())
torch.cuda.set_device(hvd.local_rank())
rank = hvd.rank()
LOGGER.info("device: {} n_gpu: {}, rank: {}, 16-bits training: {}".format(device, n_gpu, hvd.rank(), opts.fp16))
if hvd.rank() != 0:
LOGGER.disabled = True
hps_file = f'{opts.output_dir}/log/hps.json'
model_opts = Struct(load_json(hps_file))
model_config = f'{opts.output_dir}/log/model_config.json'
# load DBs and image dirs
video_ids = get_video_ids(opts.query_txt_db)
if opts.task != "didemo_video_only":
video_db = load_video_sub_dataset(opts.vfeat_db, opts.sub_txt_db, model_opts.vfeat_interval, model_opts)
else:
txt_meta = load_json(os.path.join(opts.query_txt_db, "meta.json"))
video_db = load_video_only_dataset(opts.vfeat_db, txt_meta, model_opts.vfeat_interval, model_opts)
assert opts.split in opts.query_txt_db
q_txt_db = QueryTokLmdb(opts.query_txt_db, -1)
if opts.task != "didemo_video_only":
inf_dataset = VcmrFullEvalDataset
else:
inf_dataset = VcmrVideoOnlyFullEvalDataset
eval_dataset = inf_dataset(video_ids, video_db, q_txt_db, distributed=model_opts.distributed_eval)
# Prepare model
if exists(opts.checkpoint):
ckpt_file = opts.checkpoint
else:
ckpt_file = f'{opts.output_dir}/ckpt/model_step_{opts.checkpoint}.pt'
checkpoint = torch.load(ckpt_file)
img_pos_embed_weight_key = ("v_encoder.f_encoder.img_embeddings.position_embeddings.weight")
assert img_pos_embed_weight_key in checkpoint
max_frm_seq_len = len(checkpoint[img_pos_embed_weight_key])
model = HeroForVcmr.from_pretrained(
model_config,
state_dict=checkpoint,
vfeat_dim=VFEAT_DIM,
max_frm_seq_len=max_frm_seq_len,
lw_neg_ctx=model_opts.lw_neg_ctx,
lw_neg_q=model_opts.lw_neg_q, lw_st_ed=0,
ranking_loss_type=model_opts.ranking_loss_type,
use_hard_negative=False,
hard_pool_size=model_opts.hard_pool_size,
margin=model_opts.margin,
use_all_neg=model_opts.use_all_neg,
drop_svmr_prob=model_opts.drop_svmr_prob)
model.to(device)
if opts.fp16:
model = amp.initialize(model, enabled=opts.fp16, opt_level='O2')
eval_dataloader = DataLoader(eval_dataset, batch_size=opts.batch_size,
num_workers=opts.n_workers,
pin_memory=opts.pin_mem,
collate_fn=vcmr_full_eval_collate)
eval_dataloader = PrefetchLoader(eval_dataloader)
_, results = validate_full_vcmr(model, eval_dataloader, opts.split, opts, model_opts)
result_dir = f'{opts.output_dir}/results_{opts.split}'
if not exists(result_dir) and rank == 0:
os.makedirs(result_dir)
all_results_list = all_gather_list(results)
if hvd.rank() == 0: # save for only one time
all_results = {"video2idx": all_results_list[0]["video2idx"]}
for rank_id in range(hvd.size()):
for key, val in all_results_list[rank_id].items():
if key == "video2idx":
continue
if key not in all_results:
all_results[key] = []
all_results[key].extend(all_results_list[rank_id][key])
LOGGER.info('All results joined......')
# save_vr(all_results, f'{result_dir}/results_{opts.checkpoint}_{opts.split}_vr.json')
# save_vcmr_base_on_vr(all_results, f'{result_dir}/results_{opts.checkpoint}_{opts.split}_vcmr_base_on_vr.json')
save_vcmr(all_results, f'{result_dir}/results_{opts.checkpoint}_{opts.split}_vcmr.json')
def build_target_loaders(target, tgt_ratio, opts):
if 'vfeat_shards' in target:
sub_txt_db = SubTokLmdb(f"{opts.txt_db}/{target['sub_txt_db']}",
opts.max_clip_len)
video_db = [
load_video_sub_dataset(
f"{opts.img_db}/{target['vfeat_db']}/{shard}", sub_txt_db,
target['vfeat_interval'], opts)
for shard in target['vfeat_shards']
]
else:
video_db = load_video_sub_dataset(
f"{opts.img_db}/{target['vfeat_db']}",
f"{opts.txt_db}/{target['sub_txt_db']}",
target['vfeat_interval'], opts)
train_loaders = {}
val_loaders = {}
for split in target['splits']:
if 'ratio' not in split:
split['ratio'] = [1] * len(split['tasks'])
assert len(split['tasks']) == len(split['ratio'])
for task, r in zip(split['tasks'], split['ratio']):
name = f"{task}_{target['name']}_{split['name']}"
LOGGER.info(f'loading {name} ...')
ratio = tgt_ratio * r
if isinstance(video_db, list):
all_train_ids = [
json.load(open(f"{opts.txt_db}/{ids}"))
for ids in split['train_idx']
]
else:
train_ids = json.load(
open(f"{opts.txt_db}/{split['train_idx']}"))
val_ids = json.load(open(f"{opts.txt_db}/{split['val_idx']}"))
if task == 'mlm':
if isinstance(video_db, list):
train_dset = ConcatDataset([
VideoMlmDataset(ids, vid_db, opts.mask_prob,
sub_ctx_len=opts.sub_ctx_len)
for ids, vid_db in zip(all_train_ids, video_db)
])
val_dset = VideoMlmDataset(
val_ids, video_db[0], opts.mask_prob,
sub_ctx_len=opts.sub_ctx_len)
else:
train_dset = VideoMlmDataset(
train_ids, video_db, opts.mask_prob,
sub_ctx_len=opts.sub_ctx_len)
val_dset = VideoMlmDataset(
val_ids, video_db, opts.mask_prob,
sub_ctx_len=opts.sub_ctx_len)
train_collate = mlm_collate
val_collate = mlm_collate
elif task == 'mfm-nce' or task == 'mffr':
if isinstance(video_db, list):
train_dset = ConcatDataset([
MfmDataset(ids, vid_db, opts.mask_prob)
for ids, vid_db in zip(all_train_ids, video_db)
])
val_dset = MfmDataset(val_ids, video_db[0], opts.mask_prob)
else:
train_dset = MfmDataset(train_ids, video_db,
opts.mask_prob)
val_dset = MfmDataset(val_ids, video_db, opts.mask_prob)
train_collate = mfm_collate
val_collate = mfm_collate
elif task == 'fom':
if isinstance(video_db, list):
train_dset = ConcatDataset([
FomDataset(ids, vid_db, opts.mask_prob)
for ids, vid_db in zip(all_train_ids, video_db)
])
val_dset = FomEvalDataset(val_ids, video_db[0],
opts.mask_prob)
else:
train_dset = FomDataset(train_ids, video_db,
opts.mask_prob)
val_dset = FomEvalDataset(val_ids, video_db,
opts.mask_prob)
train_collate = fom_collate
val_collate = fom_eval_collate
elif task == 'vsm':
if isinstance(video_db, list):
train_dset = ConcatDataset([
VsmDataset(ids, vid_db, sub_ctx_len=opts.sub_ctx_len)
for ids, vid_db in zip(all_train_ids, video_db)
])
val_dset = VsmDataset(val_ids, video_db[0],
sub_ctx_len=opts.sub_ctx_len)
else:
train_dset = VsmDataset(train_ids, video_db,
sub_ctx_len=opts.sub_ctx_len)
val_dset = VsmDataset(val_ids, video_db,
sub_ctx_len=opts.sub_ctx_len)
train_collate = vsm_collate
val_collate = vsm_collate
else:
raise ValueError(f'undefined task {task}')
train_loader = DataLoader(train_dset,
batch_size=opts.train_batch_size,
num_workers=opts.n_workers,
pin_memory=opts.pin_mem,
collate_fn=train_collate, shuffle=True)
val_loader = DataLoader(val_dset, batch_size=opts.val_batch_size,
num_workers=opts.n_workers,
pin_memory=opts.pin_mem,
collate_fn=val_collate, shuffle=False)
train_loaders[name] = (train_loader, ratio)
val_loaders[name] = PrefetchLoader(val_loader)
return train_loaders, val_loaders