VDGR/models/runner.py
2023-10-25 15:38:09 +02:00

831 lines
34 KiB
Python

import os
import os.path as osp
import json
from collections import deque
import time
import re
import shutil
import glob
import pickle
import gc
import numpy as np
import glog as log
try:
from apex import amp
except ModuleNotFoundError:
print('apex not found')
import torch
import torch.utils.data as tud
import torch.nn.functional as F
import torch.distributed as dist
from utils.data_utils import load_pickle_lines
from utils.visdial_metrics import SparseGTMetrics, NDCG, scores_to_ranks
import wandb
class Runner:
def __init__(self, config):
self.config = config
if 'rank' in config:
self.gpu_rank = config['rank']
else:
self.gpu_rank = 0
self.epoch_idx = 0
self.max_metric = 0.
self.max_metric_epoch_idx = 0
self.na_str = 'N/A'
if self.config["max_ckpt_to_keep"] > 0:
self.checkpoint_queue = deque(
[], maxlen=config["max_ckpt_to_keep"])
self.metrics_queue = deque([], maxlen=config["max_ckpt_to_keep"])
self.setup_wandb()
def setup_wandb(self):
if self.gpu_rank == 0:
print("[INFO] Set wandb logging on rank {}".format(0))
run = wandb.init(
project=self.config['wandb_project'], config=self.config, mode=self.config['wandb_mode'])
else:
run = None
self.run = run
def forward(self, batch, eval_visdial=False):
return NotImplementedError
def train(self, dataset, dataset_eval=None):
# wandb.login()
if os.path.exists(self.config['log_dir']) or self.config['loads_ckpt'] or self.config['loads_best_ckpt']:
self.load_ckpt()
if self.config['use_trainval']:
dataset.split = 'trainval'
else:
dataset.split = 'train'
batch_size = self.config['batch_size']
if self.config['parallel'] and self.config['dp_type'] != 'dp':
sampler_tr = tud.distributed.DistributedSampler(
dataset,
num_replicas=self.config['num_gpus'],
rank=self.gpu_rank
)
else:
sampler_tr = None
data_loader_tr = tud.DataLoader(
dataset=dataset,
batch_size=batch_size,
shuffle=self.config['training'] and not self.config['parallel'],
collate_fn=dataset.collate_fn,
num_workers=self.config['num_workers'],
sampler=sampler_tr
)
start_epoch_idx = self.epoch_idx
num_iter_epoch = self.config['num_iter_per_epoch']
if self.config['display']:
log.info(f'{num_iter_epoch} iter per epoch.')
# eval before training
eval_dense_at_first = self.config['train_on_dense'] and self.config['skip_mrr_eval'] and start_epoch_idx == 0
# eval before training under 2 circumstances:
# for dense finetuning, eval ndcg before the first epoch
# for mrr training, continue training and the last epoch is not evaluated
if (eval_dense_at_first or (self.config['eval_at_start'] and len(self.metrics_queue) == 0 and start_epoch_idx > 0)):
if eval_dense_at_first:
iter_now = 0
else:
iter_now = max(num_iter_epoch * start_epoch_idx, 0)
if dataset_eval is None:
dataset.split = 'val'
dataset_to_eval = dataset
else:
dataset_to_eval = dataset_eval
metrics_results = {}
metrics_to_maximize, metrics_results['val'] = self.evaluate(
dataset_to_eval, iter_now)
if eval_dense_at_first:
self.max_metric = metrics_to_maximize
self.max_metric_epoch_idx = -1
else:
if self.config['display']:
self.save_eval_results(
'val', start_epoch_idx - 1, metrics_results)
if metrics_to_maximize > self.max_metric:
self.max_metric = metrics_to_maximize
self.max_metric_epoch_idx = start_epoch_idx - 1
self.copy_best_results('val', start_epoch_idx - 1)
self.copy_best_predictions('val')
if dataset_eval is None:
if self.config['use_trainval']:
dataset.split = 'trainval'
else:
dataset.split = 'train'
num_epochs = self.config['num_epochs']
for epoch_idx in range(start_epoch_idx, num_epochs):
if self.config['parallel'] and self.config['dp_type'] != 'dp':
sampler_tr.set_epoch(epoch_idx)
self.epoch_idx = epoch_idx
if self.config['display']:
log.info(f'starting epoch {epoch_idx}')
log.info('training')
self.model.train()
num_batch = 0
next_logging_pct = .1
next_evaluating_pct = self.config["next_evaluating_pct"] + .1
start_time = time.time()
self.optimizer.zero_grad()
for batch in data_loader_tr:
if self.config['eval_before_training']:
log.info('Skipping stright to evaluation...')
break
num_batch += 1
pct = num_batch / num_iter_epoch * 100
iter_now = num_iter_epoch * epoch_idx + num_batch
output = self.forward(batch)
losses = output['losses']
# optimizer step
losses['tot_loss'] /= self.config['batch_multiply']
# debug
if self.config['debugging']:
log.info('try backward')
if self.config['dp_type'] == 'apex':
with amp.scale_loss(losses['tot_loss'], self.optimizer) as scaled_loss:
scaled_loss.backward()
else:
losses['tot_loss'].backward()
if self.config['debugging']:
log.info('backward done')
if iter_now % self.config['batch_multiply'] == 0:
self.optimizer.step()
self.optimizer.zero_grad()
self.scheduler.step()
# display and eval
if pct >= next_logging_pct:
if self.config['display']:
loss_to_print = ''
for key in losses:
if losses[key] is not None and isinstance(losses[key], torch.Tensor):
loss_to_print += f'[{key}: {losses[key].item():.4f}]'
print(
f'[{int(pct)}%][Epoch: {epoch_idx + 1}/{num_epochs}][Iter : {num_batch}/{len(data_loader_tr)}] [time: {time.time() - start_time:.2f}] {loss_to_print}'
)
next_logging_pct += self.config["next_logging_pct"]
if self.config['debugging']:
break
if pct >= next_evaluating_pct:
next_evaluating_pct += self.config["next_evaluating_pct"]
if self.run:
if self.config['train_on_dense']:
self.run.log(
{
"Train/dense_loss": losses['dense_loss'],
"Train/total_loss": losses['tot_loss'],
},
step=iter_now
)
else:
self.run.log(
{
"Train/lm_loss": losses['lm_loss'],
"Train/img_loss": losses['img_loss'],
"Train/nsp_loss": losses['nsp_loss'],
"Train/total_loss": losses['tot_loss'],
},
step=iter_now
)
lr_gnn, lr_bert = self.scheduler.get_lr()[0], self.scheduler.get_lr()[1]
self.run.log(
{
"Train/lr_gnn": lr_gnn,
"Train/lr_bert": lr_bert,
},
step=iter_now
)
del losses
# debug
torch.cuda.empty_cache()
if self.config['display']:
log.info(
f'100%,\ttime:\t{time.time() - start_time:.2f}'
)
ckpt_path = self.save_ckpt()
if not self.config['skip_visdial_eval'] and self.epoch_idx % self.config['eval_visdial_every'] == 0:
iter_now = num_iter_epoch * (epoch_idx + 1)
if dataset_eval is None:
dataset.split = 'val'
dataset_to_eval = dataset
else:
dataset_to_eval = dataset_eval
metrics_results = {}
metrics_to_maximize, metrics_results['val'] = self.evaluate(
dataset_to_eval, iter_now)
if dataset_eval is None:
if self.config['use_trainval']:
dataset.split = 'trainval'
else:
dataset.split = 'train'
if self.config['display']:
self.save_eval_results('val', epoch_idx, metrics_results)
if self.config['display']:
if metrics_to_maximize > self.max_metric:
self.max_metric = metrics_to_maximize
self.max_metric_epoch_idx = epoch_idx
self.copy_best_results('val', epoch_idx)
self.copy_best_predictions('val')
elif not self.config['parallel'] and epoch_idx - self.max_metric_epoch_idx > self.config["early_stop_epoch"]:
log.info('Early stop.')
break
if self.run:
self.run.log(
{"Val/metric_best": self.max_metric}, step=iter_now)
if self.config['parallel']:
if self.config['dp_type'] == 'dp':
gc.collect()
torch.cuda.empty_cache()
else:
dist.barrier()
log.info('Rank {} passed barrier...'.format(self.gpu_rank))
if self.config['stop_epochs'] >= 0 and epoch_idx + 1 >= self.config['stop_epochs']:
if self.config['display']:
log.info('Stop for reaching stop_epochs.')
break
def evaluate(self, dataset, training_iter=None, eval_visdial=True):
# create files to save output
if self.config['predicting']:
visdial_file_name = None
if self.config['save_score']:
visdial_file_name = osp.join(
self.config['log_dir'], f'visdial_prediction.pkl')
if osp.exists(visdial_file_name):
dialogs_predicted = load_pickle_lines(
visdial_file_name)
dialogs_predicted = [d['image_id']
for d in dialogs_predicted]
else:
dialogs_predicted = []
f_visdial = open(visdial_file_name, 'ab')
else:
visdial_file_name = osp.join(
self.config['log_dir'], f'visdial_prediction.jsonlines')
if self.config['parallel'] and self.config['dp_type'] != 'dp':
visdial_file_name = visdial_file_name.replace(
'.jsonlines', f'_{self.config["rank"]}of{self.config["num_gpus"]}.jsonlines')
if osp.exists(visdial_file_name):
dialogs_predicted_visdial = [json.loads(
line)['image_id'] for line in open(visdial_file_name)]
f_visdial = open(visdial_file_name, 'a')
else:
dialogs_predicted_visdial = []
f_visdial = open(visdial_file_name, 'w')
dialogs_predicted = dialogs_predicted_visdial
if len(dialogs_predicted) > 0:
log.info(f'Found {len(dialogs_predicted)} predicted results.')
if self.config['display']:
if visdial_file_name is not None:
log.info(
f'VisDial predictions saved to {visdial_file_name}')
elif self.config['display']:
if self.config['continue_evaluation']:
predicted_files = os.listdir(
osp.join(self.config['visdial_output_dir'], dataset.split))
dialogs_predicted = [
int(re.match(r'(\d+).npz', p).group(1)) for p in predicted_files]
else:
if osp.exists(osp.join(self.config['visdial_output_dir'], dataset.split)):
shutil.rmtree(
osp.join(self.config['visdial_output_dir'], dataset.split))
os.makedirs(
osp.join(self.config['visdial_output_dir'], dataset.split))
dialogs_predicted = []
log.info(f'Found {len(dialogs_predicted)} predicted results.')
if self.config['parallel'] and self.config['dp_type'] != 'dp':
sampler_val = tud.distributed.DistributedSampler(
dataset,
num_replicas=self.config['num_gpus'],
rank=self.gpu_rank
)
sampler_val.set_epoch(self.epoch_idx)
else:
sampler_val = None
data_loader_val = tud.DataLoader(
dataset=dataset,
batch_size=self.config['eval_batch_size'],
shuffle=False,
collate_fn=dataset.collate_fn,
num_workers=self.config['num_workers'],
sampler=sampler_val
)
self.model.eval()
with torch.no_grad():
if self.config['display']:
log.info(f'Evaluating {len(dataset)} samples')
next_logging_pct = self.config["next_logging_pct"] + .1
if self.config['parallel'] and self.config['dp_type'] == 'dp':
num_batch_tot = int(
np.ceil(len(dataset) / self.config['eval_batch_size']))
else:
num_batch_tot = int(np.ceil(
len(dataset) / (self.config['eval_batch_size'] * self.config['num_gpus'])))
num_batch = 0
if dataset.split == 'val':
num_options = self.config["num_options"]
if self.config['skip_mrr_eval']:
num_rounds = 1
else:
num_rounds = 10
elif dataset.split == 'test':
num_options = 100
num_rounds = 1
if self.gpu_rank == 0:
start_time = time.time()
for batch in data_loader_val:
num_batch += 1
# skip dialogs that have been predicted
if self.config['predicting']:
image_ids = batch['image_id'].tolist()
skip_batch = True
for image_id in image_ids:
if image_id not in dialogs_predicted:
skip_batch = False
if skip_batch:
continue
output = self.forward(
batch, eval_visdial=eval_visdial)
# visdial evaluation
if eval_visdial:
img_ids = batch['image_id'].tolist()
batch_size = len(img_ids)
if not self.config['skip_ndcg_eval']:
gt_relevance_round_id = batch['round_id'].tolist()
# [batch_size * num_rounds * num_options, 2]
nsp_scores = output['nsp_scores']
nsp_probs = F.softmax(nsp_scores, dim=1)
assert nsp_probs.shape[-1] == 2
# num_dim=2, 0 for postive, 1 for negative
nsp_probs = nsp_probs[:, 0]
nsp_probs = nsp_probs.view(
batch_size, num_rounds, num_options)
# could be predicting or evaluating
if dataset.split == 'val':
if self.config['skip_ndcg_eval']:
gt_option_inds = batch['gt_option_inds']
for b in range(batch_size):
filename = osp.join(
self.config['visdial_output_dir'], dataset.split, f'{img_ids[b]}.npz')
if not osp.exists(filename):
np.savez(
filename,
nsp_probs=nsp_probs[b].cpu().numpy(),
gt_option_inds=gt_option_inds[b].cpu().numpy()
)
else:
# [batch_size, num_rounds]
gt_option_inds = batch['gt_option_inds']
# [batch_size, num_options]
gt_relevance = batch['gt_relevance']
for b in range(batch_size):
filename = osp.join(
self.config['visdial_output_dir'], dataset.split, f'{img_ids[b]}.npz')
if not osp.exists(filename):
np.savez(filename,
nsp_probs=nsp_probs[b].cpu().numpy(),
gt_option_inds=gt_option_inds[b].cpu(
).numpy(),
gt_relevance=gt_relevance[b].cpu(
).numpy(),
gt_relevance_round_id=gt_relevance_round_id[b])
# must be predicting
if dataset.split == 'test':
if self.config['save_score']:
for b in range(batch_size):
prediction = {
"image_id": img_ids[b],
"nsp_probs": nsp_probs[b].cpu().numpy(),
"gt_relevance_round_id": gt_relevance_round_id[b]
}
pickle.dump(prediction, f_visdial)
else:
# [eval_batch_size, num_rounds, num_options]
ranks = scores_to_ranks(nsp_probs)
ranks = ranks.squeeze(1)
for b in range(batch_size):
prediction = {
"image_id": img_ids[b],
"round_id": gt_relevance_round_id[b],
"ranks": ranks[b].tolist()
}
f_visdial.write(json.dumps(prediction) + '\n')
# debug
if self.config['debugging']:
break
pct = num_batch / num_batch_tot * 100
if pct >= next_logging_pct:
if self.config['display'] and self.gpu_rank == 0:
log.info(
f'{int(pct)}%,\ttime:\t{time.time() - start_time:.2f}'
)
next_logging_pct += self.config["next_logging_pct"]
# debug
if self.config['debugging']:
break
if self.config['display'] and self.gpu_rank == 0:
pct = num_batch / num_batch_tot * 100
log.info(
f'{int(pct)}%,\ttime:\t{time.time() - start_time:.2f}'
)
if not self.config['validating']:
self.model.train()
if self.config['parallel'] and self.config['dp_type'] != 'dp':
dist.barrier()
print(f'{self.gpu_rank} passed barrier')
if self.config['predicting']:
f_visdial.close()
if not self.config['save_score']:
all_visdial_predictions = [json.loads(
line) for line in open(visdial_file_name)]
if self.config['predict_split'] == 'test' and len(all_visdial_predictions) == self.config['num_test_dialogs']:
visdial_file_name = visdial_file_name.replace(
'jsonlines', 'json')
with open(visdial_file_name, 'w') as f_visdial:
json.dump(all_visdial_predictions, f_visdial)
log.info(
f'Prediction for submisson save to {visdial_file_name}.')
return None, None
if self.config['display']:
if dataset.split == 'val' and eval_visdial:
if not self.config['skip_mrr_eval']:
sparse_metrics = SparseGTMetrics()
if not self.config['skip_ndcg_eval']:
ndcg = NDCG()
if dataset.split == 'val' and eval_visdial:
visdial_output_filenames = glob.glob(
osp.join(self.config['visdial_output_dir'], dataset.split, '*.npz'))
log.info(
f'Calculating visdial metrics for {len(visdial_output_filenames)} dialogs')
for visdial_output_filename in visdial_output_filenames:
output = np.load(visdial_output_filename)
nsp_probs = torch.from_numpy(
output['nsp_probs']).unsqueeze(0)
if not self.config['skip_ndcg_eval']:
gt_relevance = torch.from_numpy(output['gt_relevance']).unsqueeze(0)
if not self.config['skip_mrr_eval']:
gt_option_inds = torch.from_numpy(
output['gt_option_inds']).unsqueeze(0)
sparse_metrics.observe(nsp_probs, gt_option_inds)
if not self.config['skip_ndcg_eval']:
gt_relevance_round_id = output['gt_relevance_round_id']
nsp_probs_dense = nsp_probs[0, gt_relevance_round_id - 1, :].unsqueeze(0)
else:
nsp_probs_dense = nsp_probs.squeeze(0) # [1, 100]
if not self.config['skip_ndcg_eval']:
ndcg.observe(nsp_probs_dense, gt_relevance)
# visdial eval output
visdial_metrics = {}
if dataset.split == 'val' and eval_visdial:
if not self.config['skip_mrr_eval']:
visdial_metrics.update(sparse_metrics.retrieve(reset=True))
if not self.config['skip_ndcg_eval']:
visdial_metrics.update(ndcg.retrieve(reset=True))
if self.config['display']:
to_print = ''
for metric_name, metric_value in visdial_metrics.items():
if 'round' not in metric_name:
to_print += f"\n{metric_name}: {metric_value}"
if training_iter is not None:
if self.run:
self.run.log(
{'Val/' + metric_name: metric_value}, step=training_iter)
log.info(to_print)
if self.config['metrics_to_maximize'] in visdial_metrics:
metrics_to_maximize = visdial_metrics[self.config['metrics_to_maximize']]
else:
metrics_to_maximize = None
torch.cuda.empty_cache()
return metrics_to_maximize, visdial_metrics
else:
torch.cuda.empty_cache()
return None, None
def save_eval_results(self, split, epoch_idx, metrics_results):
metrics_filename = osp.join(
self.config['log_dir'], f'metrics_epoch_{epoch_idx}.json')
with open(metrics_filename, 'w') as f:
json.dump(metrics_results, f)
log.info(f'Results of metrics saved to {metrics_filename}')
if self.config["max_ckpt_to_keep"] > 0:
if len(self.metrics_queue) == self.metrics_queue.maxlen:
todel = self.metrics_queue.popleft()
os.remove(todel)
self.metrics_queue.append(metrics_filename)
if epoch_idx == 'best':
self.copy_best_predictions(split)
def copy_best_results(self, split, epoch_idx):
to_print = 'Copy '
if not self.config['skip_saving_ckpt']:
ckpt_path = osp.join(
self.config['log_dir'], f'epoch_{epoch_idx}.ckpt')
best_ckpt_path = ckpt_path.replace(
f'{epoch_idx}.ckpt', 'best.ckpt')
shutil.copyfile(ckpt_path, best_ckpt_path)
to_print += best_ckpt_path + ' '
metrics_filename = osp.join(
self.config['log_dir'], f'metrics_epoch_{epoch_idx}.json')
best_metric_filename = metrics_filename.replace(
f'{epoch_idx}.json', 'best.json')
shutil.copyfile(metrics_filename, best_metric_filename)
to_print += best_metric_filename + ' '
log.info(to_print)
def copy_best_predictions(self, split):
to_print = 'Copy '
visdial_output_dir = osp.join(self.config['visdial_output_dir'], split)
if osp.exists(visdial_output_dir):
dir_best = visdial_output_dir.replace('output', 'output_best')
if osp.exists(dir_best):
shutil.rmtree(dir_best)
shutil.copytree(visdial_output_dir, dir_best)
to_print += dir_best + ' '
log.info(to_print)
def get_ckpt(self):
ckpt = {
'epoch_idx': self.epoch_idx,
'max_metric': self.max_metric,
'seed': self.config['random_seed'],
'optimizer': self.optimizer.state_dict(),
'scheduler': self.scheduler.state_dict()
}
if self.config['parallel']:
ckpt['model_state_dict'] = self.model.module.state_dict()
else:
ckpt['model_state_dict'] = self.model.state_dict()
if self.config['dp_type'] == 'apex':
ckpt['amp'] = amp.state_dict()
return ckpt
def set_ckpt(self, ckpt_dict):
if not self.config['restarts']:
self.epoch_idx = ckpt_dict.get('epoch_idx', -1) + 1
if not self.config['resets_max_metric']:
self.max_metric = ckpt_dict.get('max_metric', -1)
if self.config['parallel']:
model = self.model.module
else:
model = self.model
model_state_dict = model.state_dict()
former_dict = {
k: v for k, v in ckpt_dict['model_state_dict'].items() if k in model_state_dict}
if self.config['display']:
log.info("number of keys transferred: %d" % len(former_dict))
assert len(former_dict.keys()) > 0
model_state_dict.update(former_dict)
model.load_state_dict(model_state_dict)
if self.config['display']:
log.info('loaded model')
del model_state_dict, former_dict
if not self.config['validating'] and not (self.config['uses_new_optimizer'] or self.config['sets_new_lr']):
if 'optimizer' in ckpt_dict:
self.optimizer.load_state_dict(ckpt_dict['optimizer'])
if self.config['display']:
log.info('loaded optimizer')
if 'scheduler' in ckpt_dict:
self.scheduler.last_epcoh = ckpt_dict['epoch_idx'] * \
self.config['num_iter_per_epoch']
self.scheduler.load_state_dict(ckpt_dict['scheduler'])
if 'amp' in ckpt_dict and self.config['dp_type'] == 'apex':
amp.load_state_dict(ckpt_dict['amp'])
del ckpt_dict
torch.cuda.empty_cache()
def save_ckpt(self):
ckpt_path = f'{self.config["log_dir"]}/epoch_{self.epoch_idx}.ckpt'
log.info(f'saving checkpoint {ckpt_path}')
ckpt = self.get_ckpt()
if self.config['skip_saving_ckpt']:
return ckpt_path
torch_version = float(torch.__version__[:3])
if torch_version - 1.4 > 1e-3:
torch.save(ckpt, f=ckpt_path, _use_new_zipfile_serialization=False)
else:
torch.save(ckpt, f=ckpt_path)
del ckpt
if not (self.config['parallel'] and self.config['dp_type'] in ['ddp', 'apex']):
torch.cuda.empty_cache()
if self.config["max_ckpt_to_keep"] > 0:
if len(self.checkpoint_queue) == self.checkpoint_queue.maxlen:
todel = self.checkpoint_queue.popleft()
os.remove(todel)
self.checkpoint_queue.append(ckpt_path)
return ckpt_path
def save_ckpt_best(self):
ckpt_path = f'{self.config["log_dir"]}/epoch_best.ckpt'
log.info(f'saving checkpoint {ckpt_path}')
ckpt = self.get_ckpt()
torch.save(ckpt, f=ckpt_path)
del ckpt
return ckpt_path
def load_ckpt_best(self):
ckpt_path = f'{osp.dirname(self.config["log_dir"])}/epoch_best.ckpt'
if not osp.exists(ckpt_path):
ckpt_paths = [path for path in os.listdir(
f'{self.config["log_dir"]}/') if path.endswith('.ckpt') and 'best' not in path]
if len(ckpt_paths) == 0:
if self.config['display']:
log.info(f'No .ckpt found in {self.config["log_dir"]}')
return
def sort_func(x): return int(re.search(r"(\d+)", x).groups()[0])
ckpt_path = f'{self.config["log_dir"]}/{sorted(ckpt_paths, key=sort_func, reverse=True)[0]}'
if self.config['display']:
log.info(f'loading checkpoint {ckpt_path}')
map_location = {'cuda:0': f'cuda:{self.gpu_rank}'}
self.set_ckpt(torch.load(ckpt_path, map_location=map_location))
def load_ckpt(self, ckpt_path=None):
if not ckpt_path:
if self.config['validating'] or self.config['loads_best_ckpt']:
ckpt_path = f'{self.config["log_dir"]}/epoch_best.ckpt'
else:
ckpt_paths = [path for path in os.listdir(
f'{self.config["log_dir"]}/') if path.endswith('.ckpt') and 'best' not in path]
if len(ckpt_paths) == 0:
if self.config['display']:
log.info(f'No .ckpt found in {self.config["log_dir"]}')
return
def sort_func(x): return int(
re.search(r"(\d+)", x).groups()[0])
ckpt_path = f'{self.config["log_dir"]}/{sorted(ckpt_paths, key=sort_func, reverse=True)[0]}'
if self.config['display']:
log.info(f'loading checkpoint {ckpt_path}')
epoch_name = osp.split(ckpt_path)[1].split('.')[0]
if re.search(r"(\d+)", epoch_name):
self.checkpoint_queue.append(ckpt_path)
metrics_filename = osp.join(
self.config['log_dir'], f'metrics_{epoch_name}.json')
if osp.exists(metrics_filename):
self.metrics_queue.append(metrics_filename)
map_location = {'cuda:0': f'cuda:{self.gpu_rank}'}
self.set_ckpt(torch.load(ckpt_path, map_location=map_location))
def match_model_key(self, pretrained_dict, model_dict):
matched_dict = dict()
for key in pretrained_dict:
if key in model_dict:
matched_key = key
elif key.startswith('encoder.') and key[8:] in model_dict:
matched_key = key[8:]
elif key.startswith('module.') and key[7:] in model_dict:
matched_key = key[7:]
elif 'encoder.' + key in model_dict:
matched_key = 'encoder.' + key
elif 'module.' + key in model_dict:
matched_key = 'module.' + key
else:
# not_found.append(key)
continue
matched_dict[matched_key] = pretrained_dict[key]
not_found = ""
for k in model_dict:
if k not in matched_dict:
not_found += k + '\n'
log.info("Keys from model_dict that were not found in pretrained_dict:")
log.info(not_found)
return matched_dict
def load_pretrained_vilbert(self, start_from=None):
if start_from is not None:
self.config["start_path"] = start_from
if self.config['training'] or self.config['debugging']:
ckpt_paths = [path for path in os.listdir(
f'{self.config["log_dir"]}/') if path.endswith('.ckpt') and 'best' not in path]
if len(ckpt_paths) > 0:
if self.config['display']:
log.info('Continue training')
return
if self.config['display']:
log.info(
f'Loading pretrained VilBERT from {self.config["start_path"]}')
map_location = {'cuda:0': f'cuda:{self.gpu_rank}'}
pretrained_dict = torch.load(
self.config['start_path'], map_location=map_location)
if 'model_state_dict' in pretrained_dict:
pretrained_dict = pretrained_dict['model_state_dict']
if self.config['parallel']:
model = self.model.module
else:
model = self.model
model_dict = model.state_dict()
matched_dict = self.match_model_key(pretrained_dict, model_dict)
if self.config['display']:
log.info("number of keys transferred: %d" % len(matched_dict))
assert len(matched_dict.keys()) > 0
model_dict.update(matched_dict)
model.load_state_dict(model_dict)
del pretrained_dict, model_dict, matched_dict
if not self.config['parallel'] or self.config['dp_type'] == 'dp':
torch.cuda.empty_cache()
if self.config['display']:
log.info(f'Pretrained VilBERT loaded')