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

380 lines
17 KiB
Python

import sys
from collections import OrderedDict
import torch
from torch import nn
import torch.nn.functional as F
sys.path.append('../')
from utils.model_utils import listMLE, approxNDCGLoss, listNet, neuralNDCG, neuralNDCG_transposed
from utils.data_utils import sequence_mask
from utils.optim_utils import init_optim
from models.runner import Runner
from models.vilbert_dialog import BertForMultiModalPreTraining, BertConfig
class VDGR(nn.Module):
def __init__(self, config_path, device, use_apex=False, cache_dir=None):
super(VDGR, self).__init__()
config = BertConfig.from_json_file(config_path)
self.bert_pretrained = BertForMultiModalPreTraining.from_pretrained('bert-base-uncased', config, device, use_apex=use_apex, cache_dir=cache_dir)
self.bert_pretrained.train()
def forward(self, input_ids, image_feat, image_loc, image_edge_indices, image_edge_attributes,
question_edge_indices, question_edge_attributes, question_limits,
history_edge_indices, history_sep_indices,
sep_indices=None, sep_len=None, token_type_ids=None,
attention_mask=None, masked_lm_labels=None, next_sentence_label=None,
image_attention_mask=None, image_label=None, image_target=None):
masked_lm_loss = None
masked_img_loss = None
nsp_loss = None
seq_relationship_score = None
if next_sentence_label is not None and masked_lm_labels \
is not None and image_target is not None:
# train mode, output losses
masked_lm_loss, masked_img_loss, nsp_loss, _, _, seq_relationship_score, _ = \
self.bert_pretrained(
input_ids, image_feat, image_loc, image_edge_indices, image_edge_attributes,
question_edge_indices, question_edge_attributes, question_limits,
history_edge_indices, history_sep_indices, sep_indices=sep_indices, sep_len=sep_len, \
token_type_ids=token_type_ids, attention_mask=attention_mask, masked_lm_labels=masked_lm_labels, \
next_sentence_label=next_sentence_label, image_attention_mask=image_attention_mask,\
image_label=image_label, image_target=image_target)
else:
#inference, output scores
_, _, seq_relationship_score, _, _, _ = \
self.bert_pretrained(
input_ids, image_feat, image_loc, image_edge_indices, image_edge_attributes,
question_edge_indices, question_edge_attributes, question_limits,
history_edge_indices, history_sep_indices,
sep_indices=sep_indices, sep_len=sep_len, \
token_type_ids=token_type_ids, attention_mask=attention_mask, masked_lm_labels=masked_lm_labels, \
next_sentence_label=next_sentence_label, image_attention_mask=image_attention_mask,\
image_label=image_label, image_target=image_target)
out = (masked_lm_loss, masked_img_loss, nsp_loss, seq_relationship_score)
return out
class SparseRunner(Runner):
def __init__(self, config):
super(SparseRunner, self).__init__(config)
self.model = VDGR(
self.config['model_config'], self.config['device'],
use_apex=self.config['dp_type'] == 'apex',
cache_dir=self.config['bert_cache_dir'])
self.model.to(self.config['device'])
if not self.config['validating'] or self.config['dp_type'] == 'apex':
self.optimizer, self.scheduler = init_optim(self.model, self.config)
def forward(self, batch, eval_visdial=False):
# load data
for key in batch:
if isinstance(batch[key], torch.Tensor):
batch[key] = batch[key].to(self.config['device'])
elif isinstance(batch[key], list):
if key != 'dialog_info': # Do not send the dialog_info item to the gpu
batch[key] = [x.to(self.config['device']) for x in batch[key]]
tokens = batch['tokens']
segments = batch['segments']
sep_indices = batch['sep_indices']
mask = batch['mask']
hist_len = batch['hist_len']
image_feat = batch['image_feat']
image_loc = batch['image_loc']
image_mask = batch['image_mask']
next_sentence_labels = batch.get('next_sentence_labels', None)
image_target = batch.get('image_target', None)
image_label = batch.get('image_label', None)
# load the graph data
image_edge_indices = batch['image_edge_indices']
image_edge_attributes = batch['image_edge_attributes']
question_edge_indices = batch['question_edge_indices']
question_edge_attributes = batch['question_edge_attributes']
question_limits = batch['question_limits']
history_edge_indices = batch['history_edge_indices']
history_sep_indices = batch['history_sep_indices']
sequence_lengths = torch.gather(sep_indices, 1, hist_len.view(-1, 1)) + 1
sequence_lengths = sequence_lengths.squeeze(1)
attention_mask_lm_nsp = sequence_mask(sequence_lengths, max_len=tokens.shape[1])
sep_len = hist_len + 1
losses = OrderedDict()
if eval_visdial:
num_lines = tokens.size(0)
line_batch_size = self.config['eval_line_batch_size']
num_line_batches = num_lines // line_batch_size
if num_lines % line_batch_size > 0:
num_line_batches += 1
nsp_scores = []
for j in range(num_line_batches):
# create chunks of the original batch
chunk_range = range(j*line_batch_size, min((j+1)*line_batch_size, num_lines))
tokens_chunk = tokens[chunk_range]
segments_chunk = segments[chunk_range]
sep_indices_chunk = sep_indices[chunk_range]
mask_chunk = mask[chunk_range]
sep_len_chunk = sep_len[chunk_range]
attention_mask_lm_nsp_chunk = attention_mask_lm_nsp[chunk_range]
image_feat_chunk = image_feat[chunk_range]
image_loc_chunk = image_loc[chunk_range]
image_mask_chunk = image_mask[chunk_range]
image_edge_indices_chunk = image_edge_indices[chunk_range[0]: chunk_range[-1]+1]
image_edge_attributes_chunk = image_edge_attributes[chunk_range[0]: chunk_range[-1]+1]
question_edge_indices_chunk = question_edge_indices[chunk_range[0]: chunk_range[-1]+1]
question_edge_attributes_chunk = question_edge_attributes[chunk_range[0]: chunk_range[-1]+1]
question_limits_chunk = question_limits[chunk_range[0]: chunk_range[-1]+1]
history_edge_indices_chunk = history_edge_indices[chunk_range[0]: chunk_range[-1]+1]
history_sep_indices_chunk = history_sep_indices[chunk_range[0]: chunk_range[-1]+1]
_ , _ , _, nsp_scores_chunk = \
self.model(
tokens_chunk,
image_feat_chunk,
image_loc_chunk,
image_edge_indices_chunk,
image_edge_attributes_chunk,
question_edge_indices_chunk,
question_edge_attributes_chunk,
question_limits_chunk,
history_edge_indices_chunk,
history_sep_indices_chunk,
sep_indices=sep_indices_chunk,
sep_len=sep_len_chunk,
token_type_ids=segments_chunk,
masked_lm_labels=mask_chunk,
attention_mask=attention_mask_lm_nsp_chunk,
image_attention_mask=image_mask_chunk
)
nsp_scores.append(nsp_scores_chunk)
nsp_scores = torch.cat(nsp_scores, 0)
else:
losses['lm_loss'], losses['img_loss'], losses['nsp_loss'], nsp_scores = \
self.model(
tokens,
image_feat,
image_loc,
image_edge_indices,
image_edge_attributes,
question_edge_indices,
question_edge_attributes,
question_limits,
history_edge_indices,
history_sep_indices,
next_sentence_label=next_sentence_labels,
image_target=image_target,
image_label=image_label,
sep_indices=sep_indices,
sep_len=sep_len,
token_type_ids=segments,
masked_lm_labels=mask,
attention_mask=attention_mask_lm_nsp,
image_attention_mask=image_mask
)
losses['tot_loss'] = 0
for key in ['lm_loss', 'img_loss', 'nsp_loss']:
if key in losses and losses[key] is not None:
losses[key] = losses[key].mean()
losses['tot_loss'] += self.config[f'{key}_coeff'] * losses[key]
output = {
'losses': losses,
'nsp_scores': nsp_scores
}
return output
class DenseRunner(Runner):
def __init__(self, config):
super(DenseRunner, self).__init__(config)
self.model = VDGR(
self.config['model_config'], self.config['device'],
use_apex=self.config['dp_type'] == 'apex',
cache_dir=self.config['bert_cache_dir'])
if not(self.config['parallel'] and self.config['dp_type'] == 'dp'):
self.model.to(self.config['device'])
if self.config['dense_loss'] == 'ce':
self.dense_loss = nn.KLDivLoss(reduction='batchmean')
elif self.config['dense_loss'] == 'listmle':
self.dense_loss = listMLE
elif self.config['dense_loss'] == 'listnet':
self.dense_loss = listNet
elif self.config['dense_loss'] == 'approxndcg':
self.dense_loss = approxNDCGLoss
elif self.config['dense_loss'] == 'neural_ndcg':
self.dense_loss = neuralNDCG
elif self.config['dense_loss'] == 'neural_ndcg_transposed':
self.dense_loss = neuralNDCG_transposed
else:
raise ValueError('dense_loss must be one of ce, listmle, listnet, approxndcg, neural_ndcg, neural_ndcg_transposed')
if not self.config['validating'] or self.config['dp_type'] == 'apex':
self.optimizer, self.scheduler = init_optim(self.model, self.config)
def forward(self, batch, eval_visdial=False):
# load data
for key in batch:
if isinstance(batch[key], torch.Tensor):
batch[key] = batch[key].to(self.config['device'])
elif isinstance(batch[key], list):
if key != 'dialog_info': # Do not send the dialog_info item to the gpu
batch[key] = [x.to(self.config['device']) for x in batch[key]]
# get embedding and forward visdial
tokens = batch['tokens']
segments = batch['segments']
sep_indices = batch['sep_indices']
mask = batch['mask']
hist_len = batch['hist_len']
image_feat = batch['image_feat']
image_loc = batch['image_loc']
image_mask = batch['image_mask']
next_sentence_labels = batch.get('next_sentence_labels', None)
image_target = batch.get('image_target', None)
image_label = batch.get('image_label', None)
# load the graph data
image_edge_indices = batch['image_edge_indices']
image_edge_attributes = batch['image_edge_attributes']
question_edge_indices = batch['question_edge_indices']
question_edge_attributes = batch['question_edge_attributes']
question_limits = batch['question_limits']
history_edge_indices = batch['history_edge_indices']
assert history_edge_indices[0].size(0) == 2
history_sep_indices = batch['history_sep_indices']
sequence_lengths = torch.gather(sep_indices, 1, hist_len.view(-1, 1)) + 1
sequence_lengths = sequence_lengths.squeeze(1)
attention_mask_lm_nsp = sequence_mask(sequence_lengths, max_len=tokens.shape[1])
sep_len = hist_len + 1
losses = OrderedDict()
if eval_visdial:
num_lines = tokens.size(0)
line_batch_size = self.config['eval_line_batch_size']
num_line_batches = num_lines // line_batch_size
if num_lines % line_batch_size > 0:
num_line_batches += 1
nsp_scores = []
for j in range(num_line_batches):
# create chunks of the original batch
chunk_range = range(j*line_batch_size, min((j+1)*line_batch_size, num_lines))
tokens_chunk = tokens[chunk_range]
segments_chunk = segments[chunk_range]
sep_indices_chunk = sep_indices[chunk_range]
mask_chunk = mask[chunk_range]
sep_len_chunk = sep_len[chunk_range]
attention_mask_lm_nsp_chunk = attention_mask_lm_nsp[chunk_range]
image_feat_chunk = image_feat[chunk_range]
image_loc_chunk = image_loc[chunk_range]
image_mask_chunk = image_mask[chunk_range]
image_edge_indices_chunk = image_edge_indices[chunk_range[0]: chunk_range[-1]+1]
image_edge_attributes_chunk = image_edge_attributes[chunk_range[0]: chunk_range[-1]+1]
question_edge_indices_chunk = question_edge_indices[chunk_range[0]: chunk_range[-1]+1]
question_edge_attributes_chunk = question_edge_attributes[chunk_range[0]: chunk_range[-1]+1]
question_limits_chunk = question_limits[chunk_range[0]: chunk_range[-1]+1]
history_edge_indices_chunk = history_edge_indices[chunk_range[0]: chunk_range[-1]+1]
history_sep_indices_chunk = history_sep_indices[chunk_range[0]: chunk_range[-1]+1]
_, _, _, nsp_scores_chunk = \
self.model(
tokens_chunk,
image_feat_chunk,
image_loc_chunk,
image_edge_indices_chunk,
image_edge_attributes_chunk,
question_edge_indices_chunk,
question_edge_attributes_chunk,
question_limits_chunk,
history_edge_indices_chunk,
history_sep_indices_chunk,
sep_indices=sep_indices_chunk,
sep_len=sep_len_chunk,
token_type_ids=segments_chunk,
masked_lm_labels=mask_chunk,
attention_mask=attention_mask_lm_nsp_chunk,
image_attention_mask=image_mask_chunk
)
nsp_scores.append(nsp_scores_chunk)
nsp_scores = torch.cat(nsp_scores, 0)
else:
_, _, _, nsp_scores = \
self.model(
tokens,
image_feat,
image_loc,
image_edge_indices,
image_edge_attributes,
question_edge_indices,
question_edge_attributes,
question_limits,
history_edge_indices,
history_sep_indices,
next_sentence_label=next_sentence_labels,
image_target=image_target,
image_label=image_label,
sep_indices=sep_indices,
sep_len=sep_len,
token_type_ids=segments,
masked_lm_labels=mask,
attention_mask=attention_mask_lm_nsp,
image_attention_mask=image_mask
)
if nsp_scores is not None:
nsp_scores_output = nsp_scores.detach().clone()
if not eval_visdial:
nsp_scores = nsp_scores.view(-1, self.config['num_options_dense'], 2)
if 'next_sentence_labels' in batch and self.config['nsp_loss_coeff'] > 0:
next_sentence_labels = batch['next_sentence_labels'].to(self.config['device'])
losses['nsp_loss'] = F.cross_entropy(nsp_scores.view(-1,2), next_sentence_labels.view(-1))
else:
losses['nsp_loss'] = None
if not eval_visdial:
gt_relevance = batch['gt_relevance'].to(self.config['device'])
nsp_scores = nsp_scores[:, :, 0]
if self.config['dense_loss'] == 'ce':
losses['dense_loss'] = self.dense_loss(F.log_softmax(nsp_scores, dim=1), F.softmax(gt_relevance, dim=1))
else:
losses['dense_loss'] = self.dense_loss(nsp_scores, gt_relevance)
else:
losses['dense_loss'] = None
else:
nsp_scores_output = None
losses['nsp_loss'] = None
losses['dense_loss'] = None
losses['tot_loss'] = 0
for key in ['nsp_loss', 'dense_loss']:
if key in losses and losses[key] is not None:
losses[key] = losses[key].mean()
losses['tot_loss'] += self.config[f'{key}_coeff'] * losses[key]
output = {
'losses': losses,
'nsp_scores': nsp_scores_output
}
return output