import torch import torch.nn as nn from ..utils.rewards import get_scores, get_self_cider_scores class RewardCriterion(nn.Module): def __init__(self): super(RewardCriterion, self).__init__() def forward(self, input, seq, reward): input = input.gather(2, seq.unsqueeze(2)).squeeze(2) input = input.reshape(-1) reward = reward.reshape(-1) mask = (seq>0).to(input) mask = torch.cat([mask.new(mask.size(0), 1).fill_(1), mask[:, :-1]], 1).reshape(-1) output = - input * reward * mask output = torch.sum(output) / torch.sum(mask) return output class StructureLosses(nn.Module): """ This loss is inspired by Classical Structured Prediction Losses for Sequence to Sequence Learning (Edunov et al., 2018). """ def __init__(self, opt): super(StructureLosses, self).__init__() self.opt = opt self.loss_type = opt.structure_loss_type def forward(self, input, seq, data_gts): """ Input is either logits or log softmax """ out = {} batch_size = input.size(0)# batch_size = sample_size * seq_per_img seq_per_img = batch_size // len(data_gts) assert seq_per_img == self.opt.train_sample_n, seq_per_img mask = (seq>0).to(input) mask = torch.cat([mask.new_full((mask.size(0), 1), 1), mask[:, :-1]], 1) scores = get_scores(data_gts, seq, self.opt) scores = torch.from_numpy(scores).type_as(input).view(-1, seq_per_img) out['reward'] = scores #.mean() if self.opt.entropy_reward_weight > 0: entropy = - (F.softmax(input, dim=2) * F.log_softmax(input, dim=2)).sum(2).data entropy = (entropy * mask).sum(1) / mask.sum(1) print('entropy', entropy.mean().item()) scores = scores + self.opt.entropy_reward_weight * entropy.view(-1, seq_per_img) # rescale cost to [0,1] costs = - scores if self.loss_type == 'risk' or self.loss_type == 'softmax_margin': costs = costs - costs.min(1, keepdim=True)[0] costs = costs / costs.max(1, keepdim=True)[0] # in principle # Only risk need such rescale # margin should be alright; Let's try. # Gather input: BxTxD -> BxT input = input.gather(2, seq.unsqueeze(2)).squeeze(2) if self.loss_type == 'seqnll': # input is logsoftmax input = input * mask input = input.sum(1) / mask.sum(1) input = input.view(-1, seq_per_img) target = costs.min(1)[1] output = F.cross_entropy(input, target) elif self.loss_type == 'risk': # input is logsoftmax input = input * mask input = input.sum(1) input = input.view(-1, seq_per_img) output = (F.softmax(input.exp()) * costs).sum(1).mean() # test # avg_scores = input # probs = F.softmax(avg_scores.exp_()) # loss = (probs * costs.type_as(probs)).sum() / input.size(0) # print(output.item(), loss.item()) elif self.loss_type == 'max_margin': # input is logits input = input * mask input = input.sum(1) / mask.sum(1) input = input.view(-1, seq_per_img) _, __ = costs.min(1, keepdim=True) costs_star = _ input_star = input.gather(1, __) output = F.relu(costs - costs_star - input_star + input).max(1)[0] / 2 output = output.mean() # sanity test # avg_scores = input + costs # scores_with_high_target = avg_scores.clone() # scores_with_high_target.scatter_(1, costs.min(1)[1].view(-1, 1), 1e10) # target_and_offender_index = scores_with_high_target.sort(1, True)[1][:, 0:2] # avg_scores = avg_scores.gather(1, target_and_offender_index) # target_index = avg_scores.new_zeros(avg_scores.size(0), dtype=torch.long) # loss = F.multi_margin_loss(avg_scores, target_index, size_average=True, margin=0) # print(loss.item() * 2, output.item()) elif self.loss_type == 'multi_margin': # input is logits input = input * mask input = input.sum(1) / mask.sum(1) input = input.view(-1, seq_per_img) _, __ = costs.min(1, keepdim=True) costs_star = _ input_star = input.gather(1, __) output = F.relu(costs - costs_star - input_star + input) output = output.mean() # sanity test # avg_scores = input + costs # loss = F.multi_margin_loss(avg_scores, costs.min(1)[1], margin=0) # print(output, loss) elif self.loss_type == 'softmax_margin': # input is logsoftmax input = input * mask input = input.sum(1) / mask.sum(1) input = input.view(-1, seq_per_img) input = input + costs target = costs.min(1)[1] output = F.cross_entropy(input, target) elif self.loss_type == 'real_softmax_margin': # input is logits # This is what originally defined in Kevin's paper # The result should be equivalent to softmax_margin input = input * mask input = input.sum(1) / mask.sum(1) input = input.view(-1, seq_per_img) input = input + costs target = costs.min(1)[1] output = F.cross_entropy(input, target) elif self.loss_type == 'new_self_critical': """ A different self critical Self critical uses greedy decoding score as baseline; This setting uses the average score of the rest samples as baseline (suppose c1...cn n samples, reward1 = score1 - 1/(n-1)(score2+..+scoren) ) """ baseline = (scores.sum(1, keepdim=True) - scores) / (scores.shape[1] - 1) scores = scores - baseline # self cider used as reward to promote diversity (not working that much in this way) if getattr(self.opt, 'self_cider_reward_weight', 0) > 0: _scores = get_self_cider_scores(data_gts, seq, self.opt) _scores = torch.from_numpy(_scores).type_as(scores).view(-1, 1) _scores = _scores.expand_as(scores - 1) scores += self.opt.self_cider_reward_weight * _scores output = - input * mask * scores.view(-1, 1) output = torch.sum(output) / torch.sum(mask) out['loss'] = output return out class LanguageModelCriterion(nn.Module): def __init__(self): super(LanguageModelCriterion, self).__init__() def forward(self, input, target, mask): if target.ndim == 3: target = target.reshape(-1, target.shape[2]) mask = mask.reshape(-1, mask.shape[2]) # truncate to the same size target = target[:, :input.size(1)] mask = mask[:, :input.size(1)].to(input) output = -input.gather(2, target.unsqueeze(2)).squeeze(2) * mask # Average over each token output = torch.sum(output) / torch.sum(mask) return output class LabelSmoothing(nn.Module): "Implement label smoothing." def __init__(self, size=0, padding_idx=0, smoothing=0.0): super(LabelSmoothing, self).__init__() self.criterion = nn.KLDivLoss(size_average=False, reduce=False) # self.padding_idx = padding_idx self.confidence = 1.0 - smoothing self.smoothing = smoothing # self.size = size self.true_dist = None def forward(self, input, target, mask): if target.ndim == 3: target = target.reshape(-1, target.shape[2]) mask = mask.reshape(-1, mask.shape[2]) # truncate to the same size target = target[:, :input.size(1)] mask = mask[:, :input.size(1)] input = input.reshape(-1, input.size(-1)) target = target.reshape(-1) mask = mask.reshape(-1).to(input) # assert x.size(1) == self.size self.size = input.size(1) # true_dist = x.data.clone() true_dist = input.data.clone() # true_dist.fill_(self.smoothing / (self.size - 2)) true_dist.fill_(self.smoothing / (self.size - 1)) true_dist.scatter_(1, target.data.unsqueeze(1), self.confidence) # true_dist[:, self.padding_idx] = 0 # mask = torch.nonzero(target.data == self.padding_idx) # self.true_dist = true_dist return (self.criterion(input, true_dist).sum(1) * mask).sum() / mask.sum()