"""
This includes: LossComputeBase and the standard NMTLossCompute, and
sharded loss compute stuff.
"""
import torch
import torch.nn as nn
import torch.nn.functional as F
import onmt
from onmt.modules.sparse_losses import SparsemaxLoss
from onmt.modules.sparse_activations import LogSparsemax
from onmt.constants import ModelTask, DefaultTokens
from onmt.modules.copy_generator import collapse_copy_scores
from onmt.model_builder import load_test_model
try:
import ctranslate2
except ImportError:
pass # this is tested when importing for loading a LM
[docs]class LossCompute(nn.Module):
"""
Class for managing efficient loss computation. Handles
accumulating multiple loss computations.
Args:
criterion (:obj:`nn. loss function`) : NLLoss or customed loss
generator (:obj:`nn.Module`) :
copy_attn (bool): whether copy attention mechanism is on/off
lambda_coverage: Hyper-param to apply coverage attention if any
lambda_align: Hyper-param for alignment loss
tgt_shift_index (int): 1 for NMT, 0 for LM
vocab: target vocab (for copy attention score calculation)
module that maps the output of the decoder to a
distribution over the target vocabulary.
lm_generator (:obj:`ctranslate2.Generator`): LM Generator
lm_prior_lambda (float): weight of LM model in loss
lm_prior_tau (float): scaler for LM loss
"""
def __init__(
self,
criterion,
generator,
copy_attn=False,
lambda_coverage=0.0,
lambda_align=0.0,
tgt_shift_index=1,
vocab=None,
lm_generator=None,
lm_prior_lambda=None,
lm_prior_tau=None,
lm_prior_model=None,
):
super(LossCompute, self).__init__()
self.criterion = criterion
self.generator = generator
self.lambda_coverage = lambda_coverage
self.lambda_align = lambda_align
self.tgt_shift_index = tgt_shift_index
self.copy_attn = copy_attn
self.vocab = vocab # target vocab for copy_attn need
self.lm_generator = lm_generator
self.lm_prior_lambda = lm_prior_lambda
self.lm_prior_tau = lm_prior_tau
self.lm_prior_model = lm_prior_model
[docs] @classmethod
def from_opts(cls, opt, model, vocab, train=True):
"""
Returns a subclass which wraps around an nn.Module subclass
(such as nn.NLLLoss) which defines the loss criterion. The LossCompute
object passes relevant data to a Statistics object which handles
training/validation logging.
The Criterion and LossCompute options are triggered by opt settings.
"""
device = torch.device("cuda" if onmt.utils.misc.use_gpu(opt) else "cpu")
padding_idx = vocab[DefaultTokens.PAD]
unk_idx = vocab[DefaultTokens.UNK]
if opt.lambda_coverage != 0:
assert opt.coverage_attn, (
"--coverage_attn needs to be set in "
"order to use --lambda_coverage != 0"
)
tgt_shift_idx = 1 if opt.model_task == ModelTask.SEQ2SEQ else 0
if opt.copy_attn:
criterion = onmt.modules.CopyGeneratorLoss(
len(vocab),
opt.copy_attn_force,
unk_index=unk_idx,
ignore_index=padding_idx,
)
else:
if opt.generator_function == "sparsemax":
criterion = SparsemaxLoss(ignore_index=padding_idx, reduction="sum")
else:
criterion = nn.CrossEntropyLoss(
ignore_index=padding_idx,
reduction="sum",
label_smoothing=opt.label_smoothing,
)
lm_prior_lambda = opt.lm_prior_lambda
lm_prior_tau = opt.lm_prior_tau
if opt.lm_prior_model:
if opt.lm_prior_model[-3:] == ".pt":
opt.gpu = 0
opt.fp32 = False
opt.int8 = False
_, lm_prior_model, lm_model_opt = load_test_model(
opt, model_path=opt.lm_prior_model
)
lm_prior_model.to(torch.device("cuda", opt.gpu))
lm_prior_model.eval()
lm_generator = None
else:
lm_prior_model = None
try:
import ctranslate2
lm_generator = ctranslate2.Generator(
opt.lm_prior_model, device="cuda", compute_type="float16"
)
except ImportError:
raise ImportError("Could not import ctranslate2")
else:
lm_generator = None
lm_prior_model = None
compute = cls(
criterion,
model.generator,
copy_attn=opt.copy_attn,
lambda_coverage=opt.lambda_coverage,
lambda_align=opt.lambda_align,
tgt_shift_index=tgt_shift_idx,
vocab=vocab,
lm_generator=lm_generator,
lm_prior_lambda=lm_prior_lambda,
lm_prior_tau=lm_prior_tau,
lm_prior_model=lm_prior_model,
)
compute.to(device)
return compute
@property
def padding_idx(self):
return self.criterion.ignore_index
def _compute_coverage_loss(self, std_attn, cov_attn, tgt):
"""compute coverage loss"""
zero_attn = torch.zeros(cov_attn.size()[1:], device=cov_attn.device)
cov_attn = torch.cat((zero_attn.unsqueeze(0), cov_attn[:-1]), 0)
covloss = torch.min(std_attn, cov_attn).sum(dim=-1).view(-1)
covloss[tgt == self.padding_idx] = 0
return covloss.sum()
def _compute_alignement_loss(self, align_head, ref_align):
"""Compute loss between 2 partial alignment matrix."""
# align_head contains value in [0, 1) presenting attn prob,
# 0 was resulted by the context attention src_pad_mask
# So, the correspand position in ref_align should also be 0
# Therefore, clip align_head to > 1e-18 should be bias free.
align_loss = -align_head.clamp(min=1e-18).log().mul(ref_align).sum()
align_loss *= self.lambda_align
return align_loss
def _compute_copy_loss(self, batch, output, target, align, attns):
"""Compute the copy attention loss.
Args:
batch: the current batch.
output: the predict output from the model.
target: the validate target to compare output with.
align:
attns: dictionary of attention distributions
`(tgt_len, batch, src_len)`
Returns:
A tuple with the loss and raw scores.
"""
scores = self.generator(
self._bottle(output), self._bottle(attns["copy"]), batch["src_map"]
)
loss = self.criterion(scores, align, target).sum()
return loss, scores
def _compute_lm_loss_ct2(self, output, target):
"""
Compute the loss between MT output and LM output
https://github.com/cbaziotis/lm-prior-for-nmt/blob/master
/fairseq_extension/user/lm_prior/lm_prior.py#L131-L133
"""
# rescale with tau (temperature) and apply the log_softmax.
scores = self.generator(self._bottle(output)) / self.lm_prior_tau
scores = F.log_softmax(scores.to(torch.float32), dim=-1)
src = target.detach().clone()
src[src == self.vocab[DefaultTokens.EOS]] = self.padding_idx
src = src[:, :-1, :]
src_len = src[:, :, 0].ne(self.padding_idx).sum(1)
# ct2 expects src with lengths without padding
lm_scores = self.lm_generator.forward_batch(
ctranslate2.StorageView.from_array(src[:, :, 0].to(torch.int32)),
ctranslate2.StorageView.from_array(src_len.to(torch.int32)),
return_log_probs=False,
)
lm_scores = torch.as_tensor(lm_scores, device=scores.device)
# again we use raw probs to rescale with tau and apply log_softmax
lm_scores = self._bottle(lm_scores) / self.lm_prior_tau
lm_scores = F.log_softmax(lm_scores.to(torch.float32), dim=-1)
lm_scores[:, self.vocab[DefaultTokens.UNK]] = -50
lm_scores[:, self.vocab[DefaultTokens.EOS]] -= 20
# lm_scores are in log space so log_target=True
lm_loss = F.kl_div(scores, lm_scores, reduction="none", log_target=True).sum(-1)
non_padding = self._bottle(output).ne(self.padding_idx)[:, 0]
lm_loss = lm_loss.masked_select(non_padding).sum()
lm_loss = lm_loss * (self.lm_prior_tau**2)
return lm_loss
def _compute_lm_loss(self, output, target):
"""
Compute the loss between MT output and LM output
https://github.com/cbaziotis/lm-prior-for-nmt/blob/master
/fairseq_extension/user/lm_prior/lm_prior.py#L131-L133
"""
# rescale with tau (temperature) and apply the log_softmax.
scores = self.generator(self._bottle(output)) / self.lm_prior_tau
scores = F.log_softmax(scores.to(torch.float32), dim=-1)
src = target.detach().clone()
src[src == self.vocab[DefaultTokens.EOS]] = self.padding_idx
src = src[:, :-1, :]
src_len = src[:, :, 0].ne(self.padding_idx).sum(1)
# ct2 expects src with lengths without padding
lm_outs, _ = self.lm_prior_model(src, None, src_len, with_align=False)
lm_scores = (
self.lm_prior_model.generator(self._bottle(lm_outs)).detach().clone()
/ self.lm_prior_tau
)
# again we use raw probs to rescale with tau and apply log_softmax
lm_scores = F.log_softmax(lm_scores.to(torch.float32), dim=-1)
lm_scores[:, self.vocab[DefaultTokens.UNK]] = -50
lm_scores[:, self.vocab[DefaultTokens.EOS]] -= 20
# lm_scores are in log space so log_target=True
lm_loss = F.kl_div(scores, lm_scores, reduction="none", log_target=True).sum(-1)
non_padding = self._bottle(output).ne(self.padding_idx)[:, 0]
lm_loss = lm_loss.masked_select(non_padding).sum()
lm_loss = lm_loss * (self.lm_prior_tau**2)
return lm_loss
def _bottle(self, _v):
return _v.view(-1, _v.size(2))
def _unbottle(self, _v, batch_size):
return _v.view(-1, batch_size, _v.size(1))
[docs] def ignore_prompt(self, batch):
"""
Mask the prompt in the target side of the batch examples in order
to set the loss of the prompt to zero.
For finetuning on specific tasks.
The end of the prompt must be indicated by `the DefaultTokens.MASK_BEFORE`
placeholder.
The masks are supposed to be properly handled by the loss criterion
(e.g. nn.CrossEntropyLoss ).
Args:
batch: The current batch.
"""
# Create a mask with zeros at prompt positions and ones at answer postions.
mask = batch["src"].squeeze(dim=2) == self.padding_idx
mask = torch.cumsum(mask.int(), 1)
mask = mask.unsqueeze(-1)
# Apply the mask on the target side.
batch["tgt"] *= mask.int()
# Put the padding token index at the prompt positions.
batch["tgt"] += self.padding_idx * (1 - mask.int())
return batch
[docs] def forward(self, batch, output, attns, trunc_start=0, trunc_size=None):
"""Compute the forward loss, supports truncated BPTT for long
sequences by taking a range in the decoder output sequence to
back propagate in.
Range is from `(trunc_start, trunc_start + trunc_size)`.
Truncation is an approximate efficiency trick to relieve the
memory required in the RNN buffers.
Args:
batch (batch) : batch of labeled examples
output (:obj:`FloatTensor`) :
output of decoder model ``(batch, tgt_len, hidden)``
attns (dict) : dictionary of attention weights
``(batch, tgt_len, src_len)``
trunc_start (int) : starting position of truncation window
trunc_size (int) : length of truncation window
Returns:
A tuple with the loss and a :obj:`onmt.utils.Statistics` instance.
"""
if trunc_size is None:
trunc_size = batch["tgt"].size(1) - trunc_start
# take into account here the tgt_shift_index (0 / 1 = LM/NMT)
trunc_range = (trunc_start + self.tgt_shift_index, trunc_start + trunc_size)
target = batch["tgt"][:, trunc_range[0] : trunc_range[1], :]
output = output[:, trunc_start : trunc_range[1], :].contiguous()
flat_tgt = target[:, :, 0].contiguous().view(-1)
if self.copy_attn:
align = (
batch["alignment"][:, trunc_range[0] : trunc_range[1]]
.contiguous()
.view(-1)
)
loss, scores = self._compute_copy_loss(
batch, output, flat_tgt, align, attns
)
scores_data = collapse_copy_scores(
self._unbottle(scores.clone(), len(batch["srclen"])),
batch,
self.vocab,
)
scores_data = self._bottle(scores_data)
# Correct target copy token instead of <unk>
# tgt[i] = align[i] + len(tgt_vocab)
# for i such that tgt[i] == 0 and align[i] != 0
target_data = flat_tgt.clone()
unk = self.criterion.unk_index
correct_mask = (target_data == unk) & (align != unk)
offset_align = align[correct_mask] + len(self.vocab)
target_data[correct_mask] += offset_align
scores = scores_data
flat_tgt = target_data
else:
scores = self.generator(self._bottle(output))
if isinstance(self.criterion, SparsemaxLoss):
scores = LogSparsemax(scores.to(torch.float32), dim=-1)
loss = self.criterion(scores.to(torch.float32), flat_tgt)
if self.lambda_align != 0.0:
align_head = attns["align"]
if align_head.dtype != loss.dtype: # Fix FP16
align_head = align_head.to(loss.dtype)
align_idx = batch["align"]
batch_size, pad_tgt_size, _ = batch["tgt"].size()
_, pad_src_size, _ = batch["src"].size()
align_matrix_size = [batch_size, pad_tgt_size, pad_src_size]
ref_align = onmt.utils.make_batch_align_matrix(
align_idx, align_matrix_size, normalize=True
)
ref_align = ref_align[:, trunc_range[0] : trunc_range[1], :]
if ref_align.dtype != loss.dtype:
ref_align = ref_align.to(loss.dtype)
align_loss = self._compute_alignement_loss(
align_head=align_head, ref_align=ref_align
)
loss += align_loss
if self.lambda_coverage != 0.0:
coverage_loss = self._compute_coverage_loss(
attns["std"], attns["coverage"], flat_tgt
)
loss += coverage_loss
if self.lm_generator is not None:
lm_loss = self._compute_lm_loss_ct2(output, batch["tgt"])
loss = loss + lm_loss * self.lm_prior_lambda
if self.lm_prior_model is not None:
lm_loss = self._compute_lm_loss(output, batch["tgt"])
loss = loss + lm_loss * self.lm_prior_lambda
n_sents = len(batch["srclen"]) if trunc_start == 0 else 0
stats = self._stats(n_sents, loss.sum().item(), scores, flat_tgt)
return loss, stats
def _stats(self, bsz, loss, scores, target):
"""
Args:
loss (int): the loss computed by the loss criterion.
scores (:obj:`FloatTensor`): a score for each possible output
target (:obj:`FloatTensor`): true targets
Returns:
:obj:`onmt.utils.Statistics` : statistics for this batch.
"""
pred = scores.max(1)[1]
non_padding = target.ne(self.padding_idx)
num_correct = pred.eq(target).masked_select(non_padding).sum().item()
num_non_padding = non_padding.sum().item()
n_batchs = 1 if bsz else 0
# in the case criterion reduction is None then we need
# to sum the loss of each sentence in the batch
return onmt.utils.Statistics(
loss=loss,
n_batchs=n_batchs,
n_sents=bsz,
n_words=num_non_padding,
n_correct=num_correct,
)