chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,170 @@
|
||||
#!/usr/bin/env python3
|
||||
|
||||
# Copyright (c) Facebook, Inc. and its affiliates.
|
||||
#
|
||||
# This source code is licensed under the MIT license found in the
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
|
||||
import torch
|
||||
from examples.speech_recognition.data.replabels import pack_replabels
|
||||
from fairseq import utils
|
||||
from fairseq.criterions import FairseqCriterion, register_criterion
|
||||
|
||||
|
||||
@register_criterion("asg_loss")
|
||||
class ASGCriterion(FairseqCriterion):
|
||||
@staticmethod
|
||||
def add_args(parser):
|
||||
group = parser.add_argument_group("ASG Loss")
|
||||
group.add_argument(
|
||||
"--asg-transitions-init",
|
||||
help="initial diagonal value of transition matrix",
|
||||
type=float,
|
||||
default=0.0,
|
||||
)
|
||||
group.add_argument(
|
||||
"--max-replabel", help="maximum # of replabels", type=int, default=2
|
||||
)
|
||||
group.add_argument(
|
||||
"--linseg-updates",
|
||||
help="# of training updates to use LinSeg initialization",
|
||||
type=int,
|
||||
default=0,
|
||||
)
|
||||
group.add_argument(
|
||||
"--hide-linseg-messages",
|
||||
help="hide messages about LinSeg initialization",
|
||||
action="store_true",
|
||||
)
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
task,
|
||||
silence_token,
|
||||
asg_transitions_init,
|
||||
max_replabel,
|
||||
linseg_updates,
|
||||
hide_linseg_messages,
|
||||
):
|
||||
from flashlight.lib.sequence.criterion import ASGLoss, CriterionScaleMode
|
||||
|
||||
super().__init__(task)
|
||||
self.tgt_dict = task.target_dictionary
|
||||
self.eos = self.tgt_dict.eos()
|
||||
self.silence = (
|
||||
self.tgt_dict.index(silence_token)
|
||||
if silence_token in self.tgt_dict
|
||||
else None
|
||||
)
|
||||
self.max_replabel = max_replabel
|
||||
|
||||
num_labels = len(self.tgt_dict)
|
||||
self.asg = ASGLoss(num_labels, scale_mode=CriterionScaleMode.TARGET_SZ_SQRT)
|
||||
self.asg.trans = torch.nn.Parameter(
|
||||
asg_transitions_init * torch.eye(num_labels), requires_grad=True
|
||||
)
|
||||
|
||||
self.linseg_progress = torch.nn.Parameter(
|
||||
torch.tensor([0], dtype=torch.int), requires_grad=False
|
||||
)
|
||||
self.linseg_maximum = linseg_updates
|
||||
self.linseg_message_state = "none" if hide_linseg_messages else "start"
|
||||
|
||||
@classmethod
|
||||
def build_criterion(cls, args, task):
|
||||
return cls(
|
||||
task,
|
||||
args.silence_token,
|
||||
args.asg_transitions_init,
|
||||
args.max_replabel,
|
||||
args.linseg_updates,
|
||||
args.hide_linseg_messages,
|
||||
)
|
||||
|
||||
def linseg_step(self):
|
||||
if not self.training:
|
||||
return False
|
||||
if self.linseg_progress.item() < self.linseg_maximum:
|
||||
if self.linseg_message_state == "start":
|
||||
print("| using LinSeg to initialize ASG")
|
||||
self.linseg_message_state = "finish"
|
||||
self.linseg_progress.add_(1)
|
||||
return True
|
||||
elif self.linseg_message_state == "finish":
|
||||
print("| finished LinSeg initialization")
|
||||
self.linseg_message_state = "none"
|
||||
return False
|
||||
|
||||
def replace_eos_with_silence(self, tgt):
|
||||
if tgt[-1] != self.eos:
|
||||
return tgt
|
||||
elif self.silence is None or (len(tgt) > 1 and tgt[-2] == self.silence):
|
||||
return tgt[:-1]
|
||||
else:
|
||||
return tgt[:-1] + [self.silence]
|
||||
|
||||
def forward(self, model, sample, 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
|
||||
"""
|
||||
|
||||
net_output = model(**sample["net_input"])
|
||||
emissions = net_output["encoder_out"].transpose(0, 1).contiguous()
|
||||
B = emissions.size(0)
|
||||
T = emissions.size(1)
|
||||
device = emissions.device
|
||||
|
||||
target = torch.IntTensor(B, T)
|
||||
target_size = torch.IntTensor(B)
|
||||
using_linseg = self.linseg_step()
|
||||
|
||||
for b in range(B):
|
||||
initial_target_size = sample["target_lengths"][b].item()
|
||||
if initial_target_size == 0:
|
||||
raise ValueError("target size cannot be zero")
|
||||
|
||||
tgt = sample["target"][b, :initial_target_size].tolist()
|
||||
tgt = self.replace_eos_with_silence(tgt)
|
||||
tgt = pack_replabels(tgt, self.tgt_dict, self.max_replabel)
|
||||
tgt = tgt[:T]
|
||||
|
||||
if using_linseg:
|
||||
tgt = [tgt[t * len(tgt) // T] for t in range(T)]
|
||||
|
||||
target[b][: len(tgt)] = torch.IntTensor(tgt)
|
||||
target_size[b] = len(tgt)
|
||||
|
||||
loss = self.asg.forward(emissions, target.to(device), target_size.to(device))
|
||||
|
||||
if reduce:
|
||||
loss = torch.sum(loss)
|
||||
|
||||
sample_size = (
|
||||
sample["target"].size(0) if self.args.sentence_avg else sample["ntokens"]
|
||||
)
|
||||
logging_output = {
|
||||
"loss": utils.item(loss.data) if reduce else loss.data,
|
||||
"ntokens": sample["ntokens"],
|
||||
"nsentences": sample["target"].size(0),
|
||||
"sample_size": sample_size,
|
||||
}
|
||||
return loss, sample_size, logging_output
|
||||
|
||||
@staticmethod
|
||||
def aggregate_logging_outputs(logging_outputs):
|
||||
"""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)
|
||||
agg_output = {
|
||||
"loss": loss_sum / nsentences,
|
||||
"ntokens": ntokens,
|
||||
"nsentences": nsentences,
|
||||
"sample_size": sample_size,
|
||||
}
|
||||
return agg_output
|
||||
@@ -0,0 +1,17 @@
|
||||
import importlib
|
||||
import os
|
||||
|
||||
|
||||
# ASG loss requires flashlight bindings
|
||||
files_to_skip = set()
|
||||
try:
|
||||
import flashlight.lib.sequence.criterion
|
||||
except ImportError:
|
||||
files_to_skip.add("ASG_loss.py")
|
||||
|
||||
for file in os.listdir(os.path.dirname(__file__)):
|
||||
if file.endswith(".py") and not file.startswith("_") and file not in files_to_skip:
|
||||
criterion_name = file[: file.find(".py")]
|
||||
importlib.import_module(
|
||||
"examples.speech_recognition.criterions." + criterion_name
|
||||
)
|
||||
@@ -0,0 +1,130 @@
|
||||
# Copyright (c) Facebook, Inc. and its affiliates.
|
||||
#
|
||||
# This source code is licensed under the MIT license found in the
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
|
||||
from __future__ import absolute_import, division, print_function, unicode_literals
|
||||
|
||||
import logging
|
||||
import math
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from fairseq import utils
|
||||
from fairseq.criterions import FairseqCriterion, register_criterion
|
||||
|
||||
|
||||
@register_criterion("cross_entropy_acc")
|
||||
class CrossEntropyWithAccCriterion(FairseqCriterion):
|
||||
def __init__(self, task, sentence_avg):
|
||||
super().__init__(task)
|
||||
self.sentence_avg = sentence_avg
|
||||
|
||||
def compute_loss(self, model, net_output, target, reduction, log_probs):
|
||||
# N, T -> N * T
|
||||
target = target.view(-1)
|
||||
lprobs = model.get_normalized_probs(net_output, log_probs=log_probs)
|
||||
if not hasattr(lprobs, "batch_first"):
|
||||
logging.warning(
|
||||
"ERROR: we need to know whether "
|
||||
"batch first for the net output; "
|
||||
"you need to set batch_first attribute for the return value of "
|
||||
"model.get_normalized_probs. Now, we assume this is true, but "
|
||||
"in the future, we will raise exception instead. "
|
||||
)
|
||||
batch_first = getattr(lprobs, "batch_first", True)
|
||||
if not batch_first:
|
||||
lprobs = lprobs.transpose(0, 1)
|
||||
|
||||
# N, T, D -> N * T, D
|
||||
lprobs = lprobs.view(-1, lprobs.size(-1))
|
||||
loss = F.nll_loss(
|
||||
lprobs, target, ignore_index=self.padding_idx, reduction=reduction
|
||||
)
|
||||
return lprobs, loss
|
||||
|
||||
def get_logging_output(self, sample, target, lprobs, loss):
|
||||
target = target.view(-1)
|
||||
mask = target != self.padding_idx
|
||||
correct = torch.sum(
|
||||
lprobs.argmax(1).masked_select(mask) == target.masked_select(mask)
|
||||
)
|
||||
total = torch.sum(mask)
|
||||
sample_size = (
|
||||
sample["target"].size(0) if self.sentence_avg else sample["ntokens"]
|
||||
)
|
||||
|
||||
logging_output = {
|
||||
"loss": utils.item(loss.data), # * sample['ntokens'],
|
||||
"ntokens": sample["ntokens"],
|
||||
"nsentences": sample["target"].size(0),
|
||||
"sample_size": sample_size,
|
||||
"correct": utils.item(correct.data),
|
||||
"total": utils.item(total.data),
|
||||
"nframes": torch.sum(sample["net_input"]["src_lengths"]).item(),
|
||||
}
|
||||
|
||||
return sample_size, logging_output
|
||||
|
||||
def forward(self, model, sample, reduction="sum", log_probs=True):
|
||||
"""Computes the cross entropy with accuracy metric for the given sample.
|
||||
|
||||
This is similar to CrossEntropyCriterion in fairseq, but also
|
||||
computes accuracy metrics as part of logging
|
||||
|
||||
Args:
|
||||
logprobs (Torch.tensor) of shape N, T, D i.e.
|
||||
batchsize, timesteps, dimensions
|
||||
targets (Torch.tensor) of shape N, T i.e batchsize, timesteps
|
||||
|
||||
Returns:
|
||||
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
|
||||
|
||||
TODO:
|
||||
* Currently this Criterion will only work with LSTMEncoderModels or
|
||||
FairseqModels which have decoder, or Models which return TorchTensor
|
||||
as net_output.
|
||||
We need to make a change to support all FairseqEncoder models.
|
||||
"""
|
||||
net_output = model(**sample["net_input"])
|
||||
target = model.get_targets(sample, net_output)
|
||||
lprobs, loss = self.compute_loss(
|
||||
model, net_output, target, reduction, log_probs
|
||||
)
|
||||
sample_size, logging_output = self.get_logging_output(
|
||||
sample, target, lprobs, loss
|
||||
)
|
||||
return loss, sample_size, logging_output
|
||||
|
||||
@staticmethod
|
||||
def aggregate_logging_outputs(logging_outputs):
|
||||
"""Aggregate logging outputs from data parallel training."""
|
||||
correct_sum = sum(log.get("correct", 0) for log in logging_outputs)
|
||||
total_sum = sum(log.get("total", 0) for log in logging_outputs)
|
||||
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)
|
||||
nframes = sum(log.get("nframes", 0) for log in logging_outputs)
|
||||
agg_output = {
|
||||
"loss": loss_sum / sample_size / math.log(2) if sample_size > 0 else 0.0,
|
||||
# if args.sentence_avg, then sample_size is nsentences, then loss
|
||||
# is per-sentence loss; else sample_size is ntokens, the loss
|
||||
# becomes per-output token loss
|
||||
"ntokens": ntokens,
|
||||
"nsentences": nsentences,
|
||||
"nframes": nframes,
|
||||
"sample_size": sample_size,
|
||||
"acc": correct_sum * 100.0 / total_sum if total_sum > 0 else 0.0,
|
||||
"correct": correct_sum,
|
||||
"total": total_sum,
|
||||
# total is the number of validate tokens
|
||||
}
|
||||
if sample_size != ntokens:
|
||||
agg_output["nll_loss"] = loss_sum / ntokens / math.log(2)
|
||||
# loss: per output token loss
|
||||
# nll_loss: per sentence loss
|
||||
return agg_output
|
||||
Reference in New Issue
Block a user