|
|
|
|
|
|
|
|
|
|
|
import math |
|
|
|
import torch |
|
import torch.nn.functional as F |
|
from fairseq import metrics, utils |
|
from fairseq.criterions import FairseqCriterion, register_criterion |
|
|
|
|
|
@register_criterion("sentence_ranking") |
|
class SentenceRankingCriterion(FairseqCriterion): |
|
def __init__(self, task, ranking_head_name, save_predictions, num_classes): |
|
super().__init__(task) |
|
self.ranking_head_name = ranking_head_name |
|
if save_predictions is not None: |
|
self.prediction_h = open(save_predictions, "w") |
|
else: |
|
self.prediction_h = None |
|
self.num_classes = num_classes |
|
|
|
def __del__(self): |
|
if self.prediction_h is not None: |
|
self.prediction_h.close() |
|
|
|
@staticmethod |
|
def add_args(parser): |
|
|
|
parser.add_argument('--save-predictions', metavar='FILE', |
|
help='file to save predictions to') |
|
parser.add_argument('--ranking-head-name', |
|
default='sentence_classification_head', |
|
help='name of the ranking head to use') |
|
|
|
|
|
def forward(self, model, sample, reduce=True): |
|
"""Compute ranking 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 |
|
""" |
|
assert ( |
|
hasattr(model, "classification_heads") |
|
and self.ranking_head_name in model.classification_heads |
|
), "model must provide sentence ranking head for --criterion=sentence_ranking" |
|
|
|
scores = [] |
|
for idx in range(self.num_classes): |
|
score, _ = model( |
|
**sample["net_input{idx}".format(idx=idx + 1)], |
|
classification_head_name=self.ranking_head_name, |
|
) |
|
scores.append(score) |
|
|
|
logits = torch.cat(scores, dim=1) |
|
sample_size = logits.size(0) |
|
|
|
if "target" in sample: |
|
targets = model.get_targets(sample, [logits]).view(-1) |
|
lprobs = F.log_softmax(logits, dim=-1, dtype=torch.float32) |
|
loss = F.nll_loss(lprobs, targets, reduction="sum") |
|
else: |
|
targets = None |
|
loss = torch.tensor(0.0, requires_grad=True) |
|
|
|
if self.prediction_h is not None: |
|
preds = logits.argmax(dim=1) |
|
for i, (id, pred) in enumerate(zip(sample["id"].tolist(), preds.tolist())): |
|
if targets is not None: |
|
label = targets[i].item() |
|
print("{}\t{}\t{}".format(id, pred, label), file=self.prediction_h) |
|
else: |
|
print("{}\t{}".format(id, pred), file=self.prediction_h) |
|
|
|
logging_output = { |
|
"loss": loss.data, |
|
"ntokens": sample["ntokens"], |
|
"nsentences": sample_size, |
|
"sample_size": sample_size, |
|
} |
|
if targets is not None: |
|
logging_output["ncorrect"] = (logits.argmax(dim=1) == targets).sum() |
|
|
|
return loss, sample_size, logging_output |
|
|
|
@staticmethod |
|
def reduce_metrics(logging_outputs) -> None: |
|
"""Aggregate logging outputs from data parallel training.""" |
|
loss_sum = sum(log.get("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) |
|
|
|
metrics.log_scalar( |
|
"loss", loss_sum / sample_size / math.log(2), sample_size, round=3 |
|
) |
|
if sample_size != ntokens: |
|
metrics.log_scalar( |
|
"nll_loss", loss_sum / ntokens / math.log(2), ntokens, round=3 |
|
) |
|
|
|
if len(logging_outputs) > 0 and "ncorrect" in logging_outputs[0]: |
|
ncorrect = sum(log.get("ncorrect", 0) for log in logging_outputs) |
|
metrics.log_scalar( |
|
"accuracy", 100.0 * ncorrect / nsentences, nsentences, round=1 |
|
) |
|
|
|
@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 |
|
|
