# Copyright 2022 The OFA-Sys Team. # All rights reserved. # This source code is licensed under the Apache 2.0 license # found in the LICENSE file in the root directory. import math from dataclasses import dataclass, field from typing import Optional import torch import torch.nn.functional as F import numpy as np from fairseq import metrics, utils from fairseq.criterions import FairseqCriterion, register_criterion from fairseq.dataclass import FairseqDataclass from omegaconf import II from mapcalc import calculate_map, calculate_map_range @dataclass class AdjustLabelSmoothedCrossEntropySCSTCriterionConfig(FairseqDataclass): label_smoothing: float = field( default=0.0, metadata={"help": "epsilon for label smoothing, 0 means no label smoothing"}, ) report_accuracy: bool = field( default=False, metadata={"help": "report accuracy metric"}, ) ignore_prefix_size: int = field( default=0, metadata={"help": "Ignore first N tokens"}, ) ignore_eos: bool = field( default=False, metadata={"help": "Ignore eos token"}, ) sentence_avg: bool = II("optimization.sentence_avg") drop_worst_ratio: float = field( default=0.0, metadata={"help": "ratio for discarding bad samples"}, ) drop_worst_after: int = field( default=0, metadata={"help": "steps for discarding bad samples"}, ) use_rdrop: bool = field( default=False, metadata={"help": "use R-Drop"} ) reg_alpha: float = field( default=1.0, metadata={"help": "weight for R-Drop"} ) sample_patch_num: int = field( default=196, metadata={"help": "sample patches for v1"} ) constraint_range: Optional[str] = field( default=None, metadata={"help": "constraint range"} ) acc_thresh: Optional[float] = field( default=None, metadata={"help": "acc thresh for refcoco"} ) metric: Optional[str] = field( default='acc', metadata={"help": "metric"} ) max_area_size: Optional[float] = field( default=None, metadata={"help": "max_area_size"} ) min_area_size: Optional[float] = field( default=None, metadata={"help": "min_area_size"} ) logprob: Optional[bool] = field( default=False, metadata={"help": "maximise log prob"} ) pos_reward: Optional[float] = field( default=None, metadata={"help": "pos_reward"} ) neg_reward: Optional[float] = field( default=None, metadata={"help": "neg_reward"} ) reinforce: Optional[bool] = field( default=False, metadata={"help": "reinforce"} ) lambda_reinforce: Optional[float] = field( default=0, metadata={"help": "lambda_reinforce"} ) def construct_rdrop_sample(x): if isinstance(x, dict): for key in x: x[key] = construct_rdrop_sample(x[key]) return x elif isinstance(x, torch.Tensor): return x.repeat(2, *([1] * (x.dim()-1))) elif isinstance(x, int): return x * 2 elif isinstance(x, np.ndarray): return x.repeat(2) else: raise NotImplementedError def kl_loss(p, q): p_loss = F.kl_div(p, torch.exp(q), reduction='sum') q_loss = F.kl_div(q, torch.exp(p), reduction='sum') loss = (p_loss + q_loss) / 2 return loss def label_smoothed_nll_loss( lprobs, target, epsilon, update_num, reduce=True, drop_worst_ratio=0.0, drop_worst_after=0, use_rdrop=False, reg_alpha=1.0, constraint_masks=None, constraint_start=None, constraint_end=None ): if target.dim() == lprobs.dim() - 1: target = target.unsqueeze(-1) nll_loss = -lprobs.gather(dim=-1, index=target).squeeze(-1) if constraint_masks is not None: smooth_loss = -lprobs.masked_fill(~constraint_masks, 0).sum(dim=-1, keepdim=True).squeeze(-1) eps_i = epsilon / (constraint_masks.sum(1) - 1 + 1e-6) elif constraint_start is not None and constraint_end is not None: constraint_range = [0, 1, 2, 3] + list(range(constraint_start, constraint_end)) smooth_loss = -lprobs[:, constraint_range].sum(dim=-1, keepdim=True).squeeze(-1) eps_i = epsilon / (len(constraint_range) - 1 + 1e-6) else: smooth_loss = -lprobs.sum(dim=-1, keepdim=True).squeeze(-1) eps_i = epsilon / (lprobs.size(-1) - 1) loss = (1.0 - epsilon - eps_i) * nll_loss + eps_i * smooth_loss if drop_worst_ratio > 0 and update_num > drop_worst_after: if use_rdrop: true_batch_size = loss.size(0) // 2 _, indices = torch.topk(loss[:true_batch_size], k=int(true_batch_size * (1 - drop_worst_ratio)), largest=False) loss = torch.cat([loss[indices], loss[indices+true_batch_size]]) nll_loss = torch.cat([nll_loss[indices], nll_loss[indices+true_batch_size]]) lprobs = torch.cat([lprobs[indices], lprobs[indices+true_batch_size]]) else: loss, indices = torch.topk(loss, k=int(loss.shape[0] * (1 - drop_worst_ratio)), largest=False) nll_loss = nll_loss[indices] lprobs = lprobs[indices] ntokens = loss.numel() nll_loss = nll_loss.sum() # loss = loss.sum() if use_rdrop: true_batch_size = lprobs.size(0) // 2 p = lprobs[:true_batch_size] q = lprobs[true_batch_size:] if constraint_start is not None and constraint_end is not None: constraint_range = [0, 1, 2, 3] + list(range(constraint_start, constraint_end)) p = p[:, constraint_range] q = q[:, constraint_range] loss = loss + ((kl_loss(p, q) * reg_alpha)/loss.shape[0]) return loss, nll_loss, ntokens @register_criterion( "adjust_label_smoothed_cross_entropy_scst", dataclass=AdjustLabelSmoothedCrossEntropySCSTCriterionConfig ) class AdjustLabelSmoothedCrossEntropySCSTCriterion(FairseqCriterion): def __init__( self, task, sentence_avg, label_smoothing, ignore_prefix_size=0, ignore_eos=False, report_accuracy=False, drop_worst_ratio=0, drop_worst_after=0, use_rdrop=False, reg_alpha=1.0, sample_patch_num=196, constraint_range=None, acc_thresh=None, metric='acc', max_area_size=None, min_area_size=None, logprob=False, pos_reward=None, neg_reward=None, reinforce=False, lambda_reinforce=0, ): super().__init__(task) self.sentence_avg = sentence_avg self.eps = label_smoothing self.ignore_prefix_size = ignore_prefix_size self.ignore_eos = ignore_eos self.report_accuracy = report_accuracy self.drop_worst_ratio = drop_worst_ratio self.drop_worst_after = drop_worst_after self.use_rdrop = use_rdrop self.reg_alpha = reg_alpha self.sample_patch_num = sample_patch_num 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) self.acc_thresh = acc_thresh self.metric = metric self.min_area_size = min_area_size self.max_area_size = max_area_size self.logprob = logprob self.pos_reward = pos_reward self.neg_reward = neg_reward self.reinforce = reinforce self.lambda_reinforce = lambda_reinforce def get_generator_out(self, model, sample): model.eval() with torch.no_grad(): self.task.scst_generator.model.eval() gen_out = self.task.scst_generator.generate([model], sample) hyps, refs = [], [] for i in range(len(gen_out)): hyps.append(gen_out[i][0]["tokens"][:-1] - len(self.task.src_dict) + self.task.cfg.num_bins) refs.append(sample["target"][i][:-1] - len(self.task.src_dict) + self.task.cfg.num_bins) return torch.stack(hyps, dim=0), torch.stack(refs, dim=0) def _calculate_map_score(self, hyps, refs, thresh=0.5): ground_truth = { 'boxes': refs.cpu().numpy().tolist(), 'labels': [1 for i in range(refs.shape[0])] } result_dict = { 'boxes': hyps.cpu().numpy().tolist(), 'labels': [1 for i in range(hyps.shape[0])], } score = calculate_map(ground_truth, result_dict, thresh) score = torch.tensor(score).unsqueeze(0).repeat(refs.shape[0]).to(hyps.device) return score def _calculate_ap_score(self, hyps, refs, thresh=0.5, min_area_size=None, max_area_size=None): interacts = torch.cat( [torch.where(hyps[:, :2] < refs[:, :2], refs[:, :2], hyps[:, :2]), torch.where(hyps[:, 2:] < refs[:, 2:], hyps[:, 2:], refs[:, 2:])], dim=1 ) area_predictions = (hyps[:, 2] - hyps[:, 0]) * (hyps[:, 3] - hyps[:, 1]) ## x1, y1, x2, y2, x1 < x2 area_targets = (refs[:, 2] - refs[:, 0]) * (refs[:, 3] - refs[:, 1]) interacts_w = interacts[:, 2] - interacts[:, 0] interacts_h = interacts[:, 3] - interacts[:, 1] area_interacts = interacts_w * interacts_h ious = area_interacts / (area_predictions + area_targets - area_interacts + 1e-6) if max_area_size is not None and min_area_size is not None: ious = ious * (torch.logical_or(area_targets < max_area_size, area_targets > min_area_size).float()) elif min_area_size is not None: ious = ious * (area_targets > min_area_size).float() elif max_area_size is not None: ious = ious * (area_targets < max_area_size).float() if thresh is None: return ious else: return ((ious >= thresh) & (interacts_w > 0) & (interacts_h > 0)).float() def reward_step(self, sample, model): hyps, refs = self.get_generator_out(model, sample) hyps = hyps / (self.task.cfg.num_bins - 1) * self.task.cfg.max_image_size refs = refs / (self.task.cfg.num_bins - 1) * self.task.cfg.max_image_size hyps[:, ::2] /= sample['w_resize_ratios'].unsqueeze(1) hyps[:, 1::2] /= sample['h_resize_ratios'].unsqueeze(1) refs[:, ::2] /= sample['w_resize_ratios'].unsqueeze(1) refs[:, 1::2] /= sample['h_resize_ratios'].unsqueeze(1) # scores = self._calculate_ap_score(hyps, refs) if self.metric == 'acc': scores = self._calculate_ap_score(hyps, sample['region_coords'].float(), thresh=self.acc_thresh, min_area_size=self.min_area_size, max_area_size=self.max_area_size) elif self.metric == 'map': scores = self._calculate_map_score(hyps, sample['region_coords'].float(), thresh=self.acc_thresh) else: raise NotImplemented # logging_output["_score_sum"] = scores.sum().item() # logging_output["_score_cnt"] = scores.size(0) if self.pos_reward: scores = torch.where(scores > 0, self.pos_reward, scores) if self.neg_reward: scores = torch.where(scores == 0, self.neg_reward, scores) return scores 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 """ if isinstance(sample, list): if self.sample_patch_num > 0: sample[0]['net_input']['sample_patch_num'] = self.sample_patch_num # change to support len(samples) > 2 loss_v1, sample_size_v1, logging_output_v1 = self.forward(model, sample[0], update_num, reduce) loss_v2, sample_size_v2, logging_output_v2 = self.forward(model, sample[1], update_num, reduce) loss = loss_v1 / sample_size_v1 + loss_v2 / sample_size_v2 sample_size = 1 logging_output = { "loss": loss.data, "loss_v1": loss_v1.data, "loss_v2": loss_v2.data, "nll_loss": logging_output_v1["nll_loss"].data / sample_size_v1 + logging_output_v2["nll_loss"].data / sample_size_v2, "ntokens": logging_output_v1["ntokens"] + logging_output_v2["ntokens"], "nsentences": logging_output_v1["nsentences"] + logging_output_v2["nsentences"], "sample_size": 1, "sample_size_v1": sample_size_v1, "sample_size_v2": sample_size_v2, "reward": (logging_output_v1["reward"] + logging_output_v2["reward"])/2, } return loss, sample_size, logging_output if self.use_rdrop: construct_rdrop_sample(sample) ### scst reward = self.reward_step(sample, model) # shape = bs model.train() net_output = model(**sample["net_input"]) loss, nll_loss, ntokens = self.compute_loss(model, net_output, sample, update_num, reduce=reduce, reward=reward) # loss = loss*reward loss = loss.sum() sample_size = ( sample["target"].size(0) if self.sentence_avg else ntokens ) logging_output = { "loss": loss.data, "nll_loss": nll_loss.data, "ntokens": sample["ntokens"], "nsentences": sample["nsentences"], "sample_size": sample_size, "reward": reward.mean(), } if self.report_accuracy: n_correct, total = self.compute_accuracy(model, net_output, sample) logging_output["n_correct"] = utils.item(n_correct.data) logging_output["total"] = utils.item(total.data) return loss, sample_size, logging_output def get_lprobs_and_target(self, model, net_output, sample, reward=None): conf = sample['conf'][:, None, None] if 'conf' in sample and sample['conf'] is not None else 1 constraint_masks = None if "constraint_masks" in sample and sample["constraint_masks"] is not None: constraint_masks = sample["constraint_masks"] net_output[0].masked_fill_(~constraint_masks, -math.inf) 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) * conf target = model.get_targets(sample, net_output) if self.ignore_prefix_size > 0: lprobs = lprobs[:, self.ignore_prefix_size :, :].contiguous() target = target[:, self.ignore_prefix_size :].contiguous() if constraint_masks is not None: constraint_masks = constraint_masks[:, self.ignore_prefix_size :, :].contiguous() if self.ignore_eos: bsz, seq_len, embed_dim = lprobs.size() eos_indices = target.eq(self.task.tgt_dict.eos()) lprobs = lprobs[~eos_indices].reshape(bsz, seq_len-1, embed_dim) target = target[~eos_indices].reshape(bsz, seq_len-1) if constraint_masks is not None: constraint_masks = constraint_masks[~eos_indices].reshape(bsz, seq_len-1, embed_dim) if constraint_masks is not None: constraint_masks = constraint_masks.view(-1, constraint_masks.size(-1)) if reward is not None: reward = reward.unsqueeze(1).unsqueeze(1) lprobs = lprobs*reward return lprobs.view(-1, lprobs.size(-1)), target.view(-1), constraint_masks def compute_loss(self, model, net_output, sample, update_num, reduce=True, reward=None): lprobs, target, constraint_masks = self.get_lprobs_and_target(model, net_output, sample, reward=reward) if constraint_masks is not None: constraint_masks = constraint_masks[target != self.padding_idx] # print(target.shape, self.padding_idx, lprobs.shape, target, lprobs) lprobs = lprobs[target != self.padding_idx] target = target[target != self.padding_idx] loss, nll_loss, ntokens = label_smoothed_nll_loss( lprobs, target, self.eps, update_num, reduce=reduce, drop_worst_ratio=self.drop_worst_ratio, drop_worst_after=self.drop_worst_after, use_rdrop=self.use_rdrop, reg_alpha=self.reg_alpha, constraint_masks=constraint_masks, constraint_start=self.constraint_start, constraint_end=self.constraint_end ) if self.logprob and self.reinforce: # print(-lprobs.max(dim=-1)[0].squeeze(-1).sum(), loss) if self.lambda_reinforce > 0: lprobs_, target_, constraint_masks_ = self.get_lprobs_and_target(model, net_output, sample, reward=None) loss_, _, ntokens = label_smoothed_nll_loss( lprobs_, target_, self.eps, update_num, reduce=reduce, drop_worst_ratio=self.drop_worst_ratio, drop_worst_after=self.drop_worst_after, use_rdrop=self.use_rdrop, reg_alpha=self.reg_alpha, constraint_masks=constraint_masks_, constraint_start=self.constraint_start, constraint_end=self.constraint_end ) # print(-lprobs.max(dim=-1)[0].squeeze(-1).sum(), loss_) # loss = -lprobs.max(dim=-1)[0].squeeze(-1).sum()*self.lambda_reinforce + loss_ loss = loss*self.lambda_reinforce + loss_ # only supervised with more weights via reward else: loss = -lprobs.max(dim=-1)[0].squeeze(-1).sum() elif self.logprob: loss = nll_loss return loss, nll_loss, ntokens def compute_accuracy(self, model, net_output, sample): lprobs, target = self.get_lprobs_and_target(model, net_output, sample) mask = target.ne(self.padding_idx) n_correct = torch.sum( lprobs.argmax(1).masked_select(mask).eq(target.masked_select(mask)) ) total = torch.sum(mask) return n_correct, total @classmethod 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) loss_sum_v1 = sum(log.get("loss_v1", 0) for log in logging_outputs) loss_sum_v2 = sum(log.get("loss_v2", 0) for log in logging_outputs) nll_loss_sum = sum(log.get("nll_loss", 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) sample_size_v1 = sum(log.get("sample_size_v1", 0) for log in logging_outputs) sample_size_v2 = sum(log.get("sample_size_v2", 0) for log in logging_outputs) reward = sum(log.get("reward", 0) for log in logging_outputs) metrics.log_scalar( "loss", loss_sum / sample_size, sample_size, round=3 ) metrics.log_scalar( "loss_v1", loss_sum_v1 / max(sample_size_v1, 1), max(sample_size_v1, 1), round=3 ) metrics.log_scalar( "loss_v2", loss_sum_v2 / max(sample_size_v2, 1), max(sample_size_v2, 1), round=3 ) metrics.log_scalar( "nll_loss", nll_loss_sum / sample_size, ntokens, round=3 ) metrics.log_derived( "ppl", lambda meters: utils.get_perplexity(meters["nll_loss"].avg) ) 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 ) metrics.log_scalar( "sample_size_v1", sample_size_v1, 1, round=3 ) metrics.log_scalar( "sample_size_v2", sample_size_v2, 1, round=3 ) metrics.log_scalar( "reward", reward / sample_size, sample_size, round=3 ) total = utils.item(sum(log.get("total", 0) for log in logging_outputs)) if total > 0: metrics.log_scalar("total", total) n_correct = utils.item( sum(log.get("n_correct", 0) for log in logging_outputs) ) metrics.log_scalar("n_correct", n_correct) metrics.log_derived( "accuracy", lambda meters: round( meters["n_correct"].sum * 100.0 / meters["total"].sum, 3 ) if meters["total"].sum > 0 else float("nan"), ) @staticmethod 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