Spaces:
Sleeping
Sleeping
| 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() |