Spaces:
Running
Running
# Copyright (c) Facebook, Inc. and its affiliates. | |
# | |
# This source code is licensed under the MIT license found in the | |
# LICENSE file in the root directory of this source tree. | |
import math | |
import string | |
from dataclasses import dataclass, field | |
from collections import OrderedDict | |
from typing import Optional | |
import torch | |
from fairseq import metrics, utils | |
from fairseq.criterions import FairseqCriterion, register_criterion | |
from fairseq.dataclass import FairseqDataclass | |
from omegaconf import II | |
from data import data_utils | |
from utils.cider.pyciderevalcap.ciderD.ciderD import CiderD | |
def scst_loss(lprobs, target, reward, ignore_index=None, reduce=True): | |
loss = -lprobs.gather(dim=-1, index=target.unsqueeze(-1)).squeeze() * reward.unsqueeze(-1) | |
if ignore_index is not None: | |
pad_mask = target.eq(ignore_index) | |
loss.masked_fill_(pad_mask, 0.0) | |
ntokens = (~pad_mask).sum() | |
else: | |
loss = loss.squeeze(-1) | |
ntokens = target.numel() | |
if reduce: | |
loss = loss.sum() | |
return loss, ntokens | |
class ScstRewardCriterionConfig(FairseqDataclass): | |
scst_cider_cached_tokens: str = field( | |
default="coco-train-words.p", | |
metadata={"help": "path to cached cPickle file used to calculate CIDEr scores"}, | |
) | |
ignore_prefix_size: int = field( | |
default=0, | |
metadata={"help": "Ignore first N tokens"}, | |
) | |
sentence_avg: bool = II("optimization.sentence_avg") | |
constraint_range: Optional[str] = field( | |
default=None, | |
metadata={"help": "constraint range"} | |
) | |
class ScstRewardCriterion(FairseqCriterion): | |
CIDER_REWARD_WEIGHT = 1 | |
def __init__( | |
self, | |
task, | |
scst_cider_cached_tokens, | |
sentence_avg, | |
ignore_prefix_size=0, | |
constraint_range=None | |
): | |
super().__init__(task) | |
self.scst_cider_scorer = CiderD(df=scst_cider_cached_tokens) | |
self.sentence_avg = sentence_avg | |
self.ignore_prefix_size = ignore_prefix_size | |
self.transtab = str.maketrans({key: None for key in string.punctuation}) | |
self.constraint_start = None | |
self.constraint_end = None | |
if constraint_range is not None: | |
constraint_start, constraint_end = constraint_range.split(',') | |
self.constraint_start = int(constraint_start) | |
self.constraint_end = int(constraint_end) | |
def forward(self, model, sample, update_num=0, reduce=True): | |
"""Compute the loss for the given sample. | |
Returns a tuple with three elements: | |
1) the loss | |
2) the sample size, which is used as the denominator for the gradient | |
3) logging outputs to display while training | |
""" | |
loss, score, ntokens, nsentences = self.compute_loss(model, sample, reduce=reduce) | |
sample_size = ( | |
nsentences if self.sentence_avg else ntokens | |
) | |
logging_output = { | |
"loss": loss.data, | |
"score": score, | |
"ntokens": ntokens, | |
"nsentences": nsentences, | |
"sample_size": sample_size, | |
} | |
return loss, sample_size, logging_output | |
def _calculate_eval_scores(self, gen_res, gt_idx, gt_res): | |
''' | |
gen_res: generated captions, list of str | |
gt_idx: list of int, of the same length as gen_res | |
gt_res: ground truth captions, list of list of str. | |
gen_res[i] corresponds to gt_res[gt_idx[i]] | |
Each image can have multiple ground truth captions | |
''' | |
gen_res_size = len(gen_res) | |
res = OrderedDict() | |
for i in range(gen_res_size): | |
res[i] = [self._wrap_sentence(gen_res[i].strip().translate(self.transtab))] | |
gts = OrderedDict() | |
gt_res_ = [ | |
[self._wrap_sentence(gt_res[i][j].strip().translate(self.transtab)) for j in range(len(gt_res[i]))] | |
for i in range(len(gt_res)) | |
] | |
for i in range(gen_res_size): | |
gts[i] = gt_res_[gt_idx[i]] | |
res_ = [{'image_id':i, 'caption': res[i]} for i in range(len(res))] | |
_, batch_cider_scores = self.scst_cider_scorer.compute_score(gts, res_) | |
scores = self.CIDER_REWARD_WEIGHT * batch_cider_scores | |
return scores | |
def _wrap_sentence(self, s): | |
# ensure the sentence ends with <eos> token | |
# in order to keep consisitent with cider_cached_tokens | |
r = s.strip() | |
if r.endswith('.'): | |
r = r[:-1] | |
r += ' <eos>' | |
return r | |
def get_generator_out(self, model, sample): | |
def decode(toks): | |
hypo = toks.int().cpu() | |
hypo_str = self.task.tgt_dict.string(hypo) | |
hypo_str = self.task.bpe.decode(hypo_str).strip() | |
return hypo, hypo_str | |
model.eval() | |
with torch.no_grad(): | |
self.task.scst_generator.model.eval() | |
gen_out = self.task.scst_generator.generate([model], sample) | |
gen_target = [] | |
gen_res = [] | |
gt_res = [] | |
for i in range(len(gen_out)): | |
for j in range(len(gen_out[i])): | |
hypo, hypo_str = decode(gen_out[i][j]["tokens"]) | |
gen_target.append(hypo) | |
gen_res.append(hypo_str) | |
gt_res.append( | |
decode(utils.strip_pad(sample["target"][i], self.padding_idx))[1].split('&&') | |
) | |
return gen_target, gen_res, gt_res | |
def get_reward_and_scores(self, gen_res, gt_res, device): | |
batch_size = len(gt_res) | |
gen_res_size = len(gen_res) | |
seq_per_img = gen_res_size // batch_size | |
gt_idx = [i // seq_per_img for i in range(gen_res_size)] | |
scores = self._calculate_eval_scores(gen_res, gt_idx, gt_res) | |
sc_ = scores.reshape(batch_size, seq_per_img) | |
baseline = (sc_.sum(1, keepdims=True) - sc_) / (sc_.shape[1] - 1) | |
# sample - baseline | |
reward = scores.reshape(batch_size, seq_per_img) | |
reward = reward - baseline | |
reward = reward.reshape(gen_res_size) | |
reward = torch.as_tensor(reward, device=device, dtype=torch.float64) | |
return reward, scores | |
def get_net_output(self, model, sample, gen_target): | |
def merge(sample_list, eos=self.task.tgt_dict.eos(), move_eos_to_beginning=False): | |
return data_utils.collate_tokens( | |
sample_list, | |
pad_idx=self.padding_idx, | |
eos_idx=eos, | |
left_pad=False, | |
move_eos_to_beginning=move_eos_to_beginning, | |
) | |
batch_size = len(sample["target"]) | |
gen_target_size = len(gen_target) | |
seq_per_img = gen_target_size // batch_size | |
model.train() | |
sample_src_tokens = torch.repeat_interleave( | |
sample['net_input']['src_tokens'], seq_per_img, dim=0 | |
) | |
sample_src_lengths = torch.repeat_interleave( | |
sample['net_input']['src_lengths'], seq_per_img, dim=0 | |
) | |
sample_patch_images = torch.repeat_interleave( | |
sample['net_input']['patch_images'], seq_per_img, dim=0 | |
) | |
sample_patch_masks = torch.repeat_interleave( | |
sample['net_input']['patch_masks'], seq_per_img, dim=0 | |
) | |
gen_prev_output_tokens = torch.as_tensor( | |
merge(gen_target, eos=self.task.tgt_dict.bos(), move_eos_to_beginning=True), | |
device=sample["target"].device, dtype=torch.int64 | |
) | |
gen_target_tokens = torch.as_tensor( | |
merge(gen_target), device=sample["target"].device, dtype=torch.int64 | |
) | |
net_output = model( | |
src_tokens=sample_src_tokens, src_lengths=sample_src_lengths, | |
patch_images=sample_patch_images, patch_masks=sample_patch_masks, | |
prev_output_tokens=gen_prev_output_tokens | |
) | |
return net_output, gen_target_tokens | |
def get_lprobs_and_target(self, model, net_output, gen_target): | |
if self.constraint_start is not None and self.constraint_end is not None: | |
net_output[0][:, :, 4:self.constraint_start] = -math.inf | |
net_output[0][:, :, self.constraint_end:] = -math.inf | |
lprobs = model.get_normalized_probs(net_output, log_probs=True) | |
if self.ignore_prefix_size > 0: | |
if getattr(lprobs, "batch_first", False): | |
lprobs = lprobs[:, self.ignore_prefix_size :, :].contiguous() | |
gen_target = gen_target[:, self.ignore_prefix_size :].contiguous() | |
else: | |
lprobs = lprobs[self.ignore_prefix_size :, :, :].contiguous() | |
gen_target = gen_target[self.ignore_prefix_size :, :].contiguous() | |
return lprobs, gen_target | |
def compute_loss(self, model, sample, reduce=True): | |
gen_target, gen_res, gt_res = self.get_generator_out(model, sample) | |
reward, scores = self.get_reward_and_scores(gen_res, gt_res, device=sample["target"].device) | |
net_output, gen_target_tokens = self.get_net_output(model, sample, gen_target) | |
gen_lprobs, gen_target_tokens = self.get_lprobs_and_target(model, net_output, gen_target_tokens) | |
loss, ntokens = scst_loss(gen_lprobs, gen_target_tokens, reward, ignore_index=self.padding_idx, reduce=reduce) | |
nsentences = gen_target_tokens.size(0) | |
return loss, scores.sum(), ntokens, nsentences | |
def reduce_metrics(cls, logging_outputs) -> None: | |
"""Aggregate logging outputs from data parallel training.""" | |
loss_sum = sum(log.get("loss", 0) for log in logging_outputs) | |
score_sum = sum(log.get("score", 0) for log in logging_outputs) | |
ntokens = sum(log.get("ntokens", 0) for log in logging_outputs) | |
nsentences = sum(log.get("nsentences", 0) for log in logging_outputs) | |
sample_size = sum(log.get("sample_size", 0) for log in logging_outputs) | |
metrics.log_scalar( | |
"loss", loss_sum / sample_size, sample_size, round=3 | |
) | |
metrics.log_scalar( | |
"score", score_sum / nsentences, nsentences, round=3 | |
) | |
metrics.log_scalar( | |
"ntokens", ntokens, 1, round=3 | |
) | |
metrics.log_scalar( | |
"nsentences", nsentences, 1, round=3 | |
) | |
metrics.log_scalar( | |
"sample_size", sample_size, 1, round=3 | |
) | |
def logging_outputs_can_be_summed() -> bool: | |
""" | |
Whether the logging outputs returned by `forward` can be summed | |
across workers prior to calling `reduce_metrics`. Setting this | |
to True will improves distributed training speed. | |
""" | |
return True | |