chore: import upstream snapshot with attribution
This commit is contained in:
@@ -0,0 +1,557 @@
|
||||
#!/usr/bin/env python3
|
||||
|
||||
import argparse
|
||||
import os
|
||||
import unittest
|
||||
from inspect import currentframe, getframeinfo
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
from examples.speech_recognition.data.data_utils import lengths_to_encoder_padding_mask
|
||||
from fairseq.data import data_utils as fairseq_data_utils
|
||||
from fairseq.data.dictionary import Dictionary
|
||||
from fairseq.models import (
|
||||
BaseFairseqModel,
|
||||
FairseqDecoder,
|
||||
FairseqEncoder,
|
||||
FairseqEncoderDecoderModel,
|
||||
FairseqEncoderModel,
|
||||
FairseqModel,
|
||||
)
|
||||
from fairseq.tasks.fairseq_task import LegacyFairseqTask
|
||||
|
||||
|
||||
DEFAULT_TEST_VOCAB_SIZE = 100
|
||||
|
||||
|
||||
# ///////////////////////////////////////////////////////////////////////////
|
||||
# utility function to setup dummy dict/task/input
|
||||
# ///////////////////////////////////////////////////////////////////////////
|
||||
|
||||
|
||||
def get_dummy_dictionary(vocab_size=DEFAULT_TEST_VOCAB_SIZE):
|
||||
dummy_dict = Dictionary()
|
||||
# add dummy symbol to satisfy vocab size
|
||||
for id, _ in enumerate(range(vocab_size)):
|
||||
dummy_dict.add_symbol("{}".format(id), 1000)
|
||||
return dummy_dict
|
||||
|
||||
|
||||
class DummyTask(LegacyFairseqTask):
|
||||
def __init__(self, args):
|
||||
super().__init__(args)
|
||||
self.dictionary = get_dummy_dictionary()
|
||||
if getattr(self.args, "ctc", False):
|
||||
self.dictionary.add_symbol("<ctc_blank>")
|
||||
self.tgt_dict = self.dictionary
|
||||
|
||||
@property
|
||||
def target_dictionary(self):
|
||||
return self.dictionary
|
||||
|
||||
|
||||
def get_dummy_task_and_parser():
|
||||
"""
|
||||
to build a fariseq model, we need some dummy parse and task. This function
|
||||
is used to create dummy task and parser to faciliate model/criterion test
|
||||
|
||||
Note: we use FbSpeechRecognitionTask as the dummy task. You may want
|
||||
to use other task by providing another function
|
||||
"""
|
||||
parser = argparse.ArgumentParser(
|
||||
description="test_dummy_s2s_task", argument_default=argparse.SUPPRESS
|
||||
)
|
||||
DummyTask.add_args(parser)
|
||||
args = parser.parse_args([])
|
||||
task = DummyTask.setup_task(args)
|
||||
return task, parser
|
||||
|
||||
|
||||
def get_dummy_input(T=100, D=80, B=5, K=100):
|
||||
forward_input = {}
|
||||
# T max sequence length
|
||||
# D feature vector dimension
|
||||
# B batch size
|
||||
# K target dimension size
|
||||
feature = torch.randn(B, T, D)
|
||||
# this (B, T, D) layout is just a convention, you can override it by
|
||||
# write your own _prepare_forward_input function
|
||||
src_lengths = torch.from_numpy(
|
||||
np.random.randint(low=1, high=T, size=B, dtype=np.int64)
|
||||
)
|
||||
src_lengths[0] = T # make sure the maximum length matches
|
||||
prev_output_tokens = []
|
||||
for b in range(B):
|
||||
token_length = np.random.randint(low=1, high=src_lengths[b].item() + 1)
|
||||
tokens = np.random.randint(low=0, high=K, size=token_length, dtype=np.int64)
|
||||
prev_output_tokens.append(torch.from_numpy(tokens))
|
||||
|
||||
prev_output_tokens = fairseq_data_utils.collate_tokens(
|
||||
prev_output_tokens,
|
||||
pad_idx=1,
|
||||
eos_idx=2,
|
||||
left_pad=False,
|
||||
move_eos_to_beginning=False,
|
||||
)
|
||||
src_lengths, sorted_order = src_lengths.sort(descending=True)
|
||||
forward_input["src_tokens"] = feature.index_select(0, sorted_order)
|
||||
forward_input["src_lengths"] = src_lengths
|
||||
forward_input["prev_output_tokens"] = prev_output_tokens
|
||||
|
||||
return forward_input
|
||||
|
||||
|
||||
def get_dummy_encoder_output(encoder_out_shape=(100, 80, 5)):
|
||||
"""
|
||||
This only provides an example to generate dummy encoder output
|
||||
"""
|
||||
(T, B, D) = encoder_out_shape
|
||||
encoder_out = {}
|
||||
|
||||
encoder_out["encoder_out"] = torch.from_numpy(
|
||||
np.random.randn(*encoder_out_shape).astype(np.float32)
|
||||
)
|
||||
seq_lengths = torch.from_numpy(np.random.randint(low=1, high=T, size=B))
|
||||
# some dummy mask
|
||||
encoder_out["encoder_padding_mask"] = torch.arange(T).view(1, T).expand(
|
||||
B, -1
|
||||
) >= seq_lengths.view(B, 1).expand(-1, T)
|
||||
encoder_out["encoder_padding_mask"].t_()
|
||||
|
||||
# encoer_padding_mask is (T, B) tensor, with (t, b)-th element indicate
|
||||
# whether encoder_out[t, b] is valid (=0) or not (=1)
|
||||
return encoder_out
|
||||
|
||||
|
||||
def _current_postion_info():
|
||||
cf = currentframe()
|
||||
frameinfo = " (at {}:{})".format(
|
||||
os.path.basename(getframeinfo(cf).filename), cf.f_back.f_lineno
|
||||
)
|
||||
return frameinfo
|
||||
|
||||
|
||||
def check_encoder_output(encoder_output, batch_size=None):
|
||||
"""we expect encoder_output to be a dict with the following
|
||||
key/value pairs:
|
||||
- encoder_out: a Torch.Tensor
|
||||
- encoder_padding_mask: a binary Torch.Tensor
|
||||
"""
|
||||
if not isinstance(encoder_output, dict):
|
||||
msg = (
|
||||
"FairseqEncoderModel.forward(...) must be a dict" + _current_postion_info()
|
||||
)
|
||||
return False, msg
|
||||
|
||||
if "encoder_out" not in encoder_output:
|
||||
msg = (
|
||||
"FairseqEncoderModel.forward(...) must contain encoder_out"
|
||||
+ _current_postion_info()
|
||||
)
|
||||
return False, msg
|
||||
|
||||
if "encoder_padding_mask" not in encoder_output:
|
||||
msg = (
|
||||
"FairseqEncoderModel.forward(...) must contain encoder_padding_mask"
|
||||
+ _current_postion_info()
|
||||
)
|
||||
return False, msg
|
||||
|
||||
if not isinstance(encoder_output["encoder_out"], torch.Tensor):
|
||||
msg = "encoder_out must be a torch.Tensor" + _current_postion_info()
|
||||
return False, msg
|
||||
|
||||
if encoder_output["encoder_out"].dtype != torch.float32:
|
||||
msg = "encoder_out must have float32 dtype" + _current_postion_info()
|
||||
return False, msg
|
||||
|
||||
mask = encoder_output["encoder_padding_mask"]
|
||||
if mask is not None:
|
||||
if not isinstance(mask, torch.Tensor):
|
||||
msg = (
|
||||
"encoder_padding_mask must be a torch.Tensor" + _current_postion_info()
|
||||
)
|
||||
return False, msg
|
||||
if mask.dtype != torch.uint8 and (
|
||||
not hasattr(torch, "bool") or mask.dtype != torch.bool
|
||||
):
|
||||
msg = (
|
||||
"encoder_padding_mask must have dtype of uint8"
|
||||
+ _current_postion_info()
|
||||
)
|
||||
return False, msg
|
||||
|
||||
if mask.dim() != 2:
|
||||
msg = (
|
||||
"we expect encoder_padding_mask to be a 2-d tensor, in shape (T, B)"
|
||||
+ _current_postion_info()
|
||||
)
|
||||
return False, msg
|
||||
|
||||
if batch_size is not None and mask.size(1) != batch_size:
|
||||
msg = (
|
||||
"we expect encoder_padding_mask to be a 2-d tensor, with size(1)"
|
||||
+ " being the batch size"
|
||||
+ _current_postion_info()
|
||||
)
|
||||
return False, msg
|
||||
return True, None
|
||||
|
||||
|
||||
def check_decoder_output(decoder_output):
|
||||
"""we expect output from a decoder is a tuple with the following constraint:
|
||||
- the first element is a torch.Tensor
|
||||
- the second element can be anything (reserved for future use)
|
||||
"""
|
||||
if not isinstance(decoder_output, tuple):
|
||||
msg = "FariseqDecoder output must be a tuple" + _current_postion_info()
|
||||
return False, msg
|
||||
|
||||
if len(decoder_output) != 2:
|
||||
msg = "FairseqDecoder output must be 2-elem tuple" + _current_postion_info()
|
||||
return False, msg
|
||||
|
||||
if not isinstance(decoder_output[0], torch.Tensor):
|
||||
msg = (
|
||||
"FariseqDecoder output[0] must be a torch.Tensor" + _current_postion_info()
|
||||
)
|
||||
return False, msg
|
||||
|
||||
return True, None
|
||||
|
||||
|
||||
# ///////////////////////////////////////////////////////////////////////////
|
||||
# Base Test class
|
||||
# ///////////////////////////////////////////////////////////////////////////
|
||||
|
||||
|
||||
class TestBaseFairseqModelBase(unittest.TestCase):
|
||||
"""
|
||||
This class is used to facilitate writing unittest for any class derived from
|
||||
`BaseFairseqModel`.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
if cls is TestBaseFairseqModelBase:
|
||||
raise unittest.SkipTest("Skipping test case in base")
|
||||
super().setUpClass()
|
||||
|
||||
def setUpModel(self, model):
|
||||
self.assertTrue(isinstance(model, BaseFairseqModel))
|
||||
self.model = model
|
||||
|
||||
def setupInput(self):
|
||||
pass
|
||||
|
||||
def setUp(self):
|
||||
self.model = None
|
||||
self.forward_input = None
|
||||
pass
|
||||
|
||||
|
||||
class TestFairseqEncoderDecoderModelBase(TestBaseFairseqModelBase):
|
||||
"""
|
||||
base code to test FairseqEncoderDecoderModel (formally known as
|
||||
`FairseqModel`) must be derived from this base class
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
if cls is TestFairseqEncoderDecoderModelBase:
|
||||
raise unittest.SkipTest("Skipping test case in base")
|
||||
super().setUpClass()
|
||||
|
||||
def setUpModel(self, model_cls, extra_args_setters=None):
|
||||
self.assertTrue(
|
||||
issubclass(model_cls, (FairseqEncoderDecoderModel, FairseqModel)),
|
||||
msg="This class only tests for FairseqModel subclasses",
|
||||
)
|
||||
|
||||
task, parser = get_dummy_task_and_parser()
|
||||
model_cls.add_args(parser)
|
||||
|
||||
args = parser.parse_args([])
|
||||
|
||||
if extra_args_setters is not None:
|
||||
for args_setter in extra_args_setters:
|
||||
args_setter(args)
|
||||
model = model_cls.build_model(args, task)
|
||||
self.model = model
|
||||
|
||||
def setUpInput(self, input=None):
|
||||
self.forward_input = get_dummy_input() if input is None else input
|
||||
|
||||
def setUp(self):
|
||||
super().setUp()
|
||||
|
||||
def test_forward(self):
|
||||
if self.model and self.forward_input:
|
||||
forward_output = self.model.forward(**self.forward_input)
|
||||
# for FairseqEncoderDecoderModel, forward returns a tuple of two
|
||||
# elements, the first one is a Torch.Tensor
|
||||
succ, msg = check_decoder_output(forward_output)
|
||||
if not succ:
|
||||
self.assertTrue(succ, msg=msg)
|
||||
self.forward_output = forward_output
|
||||
|
||||
def test_get_normalized_probs(self):
|
||||
if self.model and self.forward_input:
|
||||
forward_output = self.model.forward(**self.forward_input)
|
||||
logprob = self.model.get_normalized_probs(forward_output, log_probs=True)
|
||||
prob = self.model.get_normalized_probs(forward_output, log_probs=False)
|
||||
|
||||
# in order for different models/criterion to play with each other
|
||||
# we need to know whether the logprob or prob output is batch_first
|
||||
# or not. We assume an additional attribute will be attached to logprob
|
||||
# or prob. If you find your code failed here, simply override
|
||||
# FairseqModel.get_normalized_probs, see example at
|
||||
# https://fburl.com/batch_first_example
|
||||
self.assertTrue(hasattr(logprob, "batch_first"))
|
||||
self.assertTrue(hasattr(prob, "batch_first"))
|
||||
|
||||
self.assertTrue(torch.is_tensor(logprob))
|
||||
self.assertTrue(torch.is_tensor(prob))
|
||||
|
||||
|
||||
class TestFairseqEncoderModelBase(TestBaseFairseqModelBase):
|
||||
"""
|
||||
base class to test FairseqEncoderModel
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
if cls is TestFairseqEncoderModelBase:
|
||||
raise unittest.SkipTest("Skipping test case in base")
|
||||
super().setUpClass()
|
||||
|
||||
def setUpModel(self, model_cls, extra_args_setters=None):
|
||||
self.assertTrue(
|
||||
issubclass(model_cls, FairseqEncoderModel),
|
||||
msg="This class is only used for testing FairseqEncoderModel",
|
||||
)
|
||||
task, parser = get_dummy_task_and_parser()
|
||||
model_cls.add_args(parser)
|
||||
args = parser.parse_args([])
|
||||
if extra_args_setters is not None:
|
||||
for args_setter in extra_args_setters:
|
||||
args_setter(args)
|
||||
|
||||
model = model_cls.build_model(args, task)
|
||||
self.model = model
|
||||
|
||||
def setUpInput(self, input=None):
|
||||
self.forward_input = get_dummy_input() if input is None else input
|
||||
# get_dummy_input() is originally for s2s, here we delete extra dict
|
||||
# items, so it can be used for EncoderModel / Encoder as well
|
||||
self.forward_input.pop("prev_output_tokens", None)
|
||||
|
||||
def setUp(self):
|
||||
super().setUp()
|
||||
|
||||
def test_forward(self):
|
||||
if self.forward_input and self.model:
|
||||
bsz = self.forward_input["src_tokens"].size(0)
|
||||
forward_output = self.model.forward(**self.forward_input)
|
||||
|
||||
# we expect forward_output to be a dict with the following
|
||||
# key/value pairs:
|
||||
# - encoder_out: a Torch.Tensor
|
||||
# - encoder_padding_mask: a binary Torch.Tensor
|
||||
succ, msg = check_encoder_output(forward_output, batch_size=bsz)
|
||||
if not succ:
|
||||
self.assertTrue(succ, msg=msg)
|
||||
self.forward_output = forward_output
|
||||
|
||||
def test_get_normalized_probs(self):
|
||||
if self.model and self.forward_input:
|
||||
forward_output = self.model.forward(**self.forward_input)
|
||||
logprob = self.model.get_normalized_probs(forward_output, log_probs=True)
|
||||
prob = self.model.get_normalized_probs(forward_output, log_probs=False)
|
||||
|
||||
# in order for different models/criterion to play with each other
|
||||
# we need to know whether the logprob or prob output is batch_first
|
||||
# or not. We assume an additional attribute will be attached to logprob
|
||||
# or prob. If you find your code failed here, simply override
|
||||
# FairseqModel.get_normalized_probs, see example at
|
||||
# https://fburl.com/batch_first_example
|
||||
self.assertTrue(hasattr(logprob, "batch_first"))
|
||||
self.assertTrue(hasattr(prob, "batch_first"))
|
||||
|
||||
self.assertTrue(torch.is_tensor(logprob))
|
||||
self.assertTrue(torch.is_tensor(prob))
|
||||
|
||||
|
||||
class TestFairseqEncoderBase(unittest.TestCase):
|
||||
"""
|
||||
base class to test FairseqEncoder
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
if cls is TestFairseqEncoderBase:
|
||||
raise unittest.SkipTest("Skipping test case in base")
|
||||
super().setUpClass()
|
||||
|
||||
def setUpEncoder(self, encoder):
|
||||
self.assertTrue(
|
||||
isinstance(encoder, FairseqEncoder),
|
||||
msg="This class is only used for test FairseqEncoder",
|
||||
)
|
||||
self.encoder = encoder
|
||||
|
||||
def setUpInput(self, input=None):
|
||||
self.forward_input = get_dummy_input() if input is None else input
|
||||
# get_dummy_input() is originally for s2s, here we delete extra dict
|
||||
# items, so it can be used for EncoderModel / Encoder as well
|
||||
self.forward_input.pop("prev_output_tokens", None)
|
||||
|
||||
def setUp(self):
|
||||
self.encoder = None
|
||||
self.forward_input = None
|
||||
|
||||
def test_forward(self):
|
||||
if self.encoder and self.forward_input:
|
||||
bsz = self.forward_input["src_tokens"].size(0)
|
||||
|
||||
forward_output = self.encoder.forward(**self.forward_input)
|
||||
succ, msg = check_encoder_output(forward_output, batch_size=bsz)
|
||||
if not succ:
|
||||
self.assertTrue(succ, msg=msg)
|
||||
self.forward_output = forward_output
|
||||
|
||||
|
||||
class TestFairseqDecoderBase(unittest.TestCase):
|
||||
"""
|
||||
base class to test FairseqDecoder
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
if cls is TestFairseqDecoderBase:
|
||||
raise unittest.SkipTest("Skipping test case in base")
|
||||
super().setUpClass()
|
||||
|
||||
def setUpDecoder(self, decoder):
|
||||
self.assertTrue(
|
||||
isinstance(decoder, FairseqDecoder),
|
||||
msg="This class is only used for test FairseqDecoder",
|
||||
)
|
||||
self.decoder = decoder
|
||||
|
||||
def setUpInput(self, input=None):
|
||||
self.forward_input = get_dummy_encoder_output() if input is None else input
|
||||
|
||||
def setUpPrevOutputTokens(self, tokens=None):
|
||||
if tokens is None:
|
||||
self.encoder_input = get_dummy_input()
|
||||
self.prev_output_tokens = self.encoder_input["prev_output_tokens"]
|
||||
else:
|
||||
self.prev_output_tokens = tokens
|
||||
|
||||
def setUp(self):
|
||||
self.decoder = None
|
||||
self.forward_input = None
|
||||
self.prev_output_tokens = None
|
||||
|
||||
def test_forward(self):
|
||||
if (
|
||||
self.decoder is not None
|
||||
and self.forward_input is not None
|
||||
and self.prev_output_tokens is not None
|
||||
):
|
||||
forward_output = self.decoder.forward(
|
||||
prev_output_tokens=self.prev_output_tokens,
|
||||
encoder_out=self.forward_input,
|
||||
)
|
||||
succ, msg = check_decoder_output(forward_output)
|
||||
if not succ:
|
||||
self.assertTrue(succ, msg=msg)
|
||||
self.forward_input = forward_output
|
||||
|
||||
|
||||
class DummyEncoderModel(FairseqEncoderModel):
|
||||
def __init__(self, encoder):
|
||||
super().__init__(encoder)
|
||||
|
||||
@classmethod
|
||||
def build_model(cls, args, task):
|
||||
return cls(DummyEncoder())
|
||||
|
||||
def get_logits(self, net_output):
|
||||
# Inverse of sigmoid to use with BinaryCrossEntropyWithLogitsCriterion as
|
||||
# F.binary_cross_entropy_with_logits combines sigmoid and CE
|
||||
return torch.log(
|
||||
torch.div(net_output["encoder_out"], 1 - net_output["encoder_out"])
|
||||
)
|
||||
|
||||
def get_normalized_probs(self, net_output, log_probs, sample=None):
|
||||
lprobs = super().get_normalized_probs(net_output, log_probs, sample=sample)
|
||||
lprobs.batch_first = True
|
||||
return lprobs
|
||||
|
||||
|
||||
class DummyEncoder(FairseqEncoder):
|
||||
def __init__(self):
|
||||
super().__init__(None)
|
||||
|
||||
def forward(self, src_tokens, src_lengths):
|
||||
mask, max_len = lengths_to_encoder_padding_mask(src_lengths)
|
||||
return {"encoder_out": src_tokens, "encoder_padding_mask": mask}
|
||||
|
||||
|
||||
class CrossEntropyCriterionTestBase(unittest.TestCase):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
if cls is CrossEntropyCriterionTestBase:
|
||||
raise unittest.SkipTest("Skipping base class test case")
|
||||
super().setUpClass()
|
||||
|
||||
def setUpArgs(self):
|
||||
args = argparse.Namespace()
|
||||
args.sentence_avg = False
|
||||
args.threshold = 0.1 # to use with BinaryCrossEntropyWithLogitsCriterion
|
||||
return args
|
||||
|
||||
def setUp(self):
|
||||
args = self.setUpArgs()
|
||||
self.model = DummyEncoderModel(encoder=DummyEncoder())
|
||||
self.criterion = self.criterion_cls.build_criterion(args, task=DummyTask(args))
|
||||
|
||||
def get_src_tokens(self, correct_prediction, aggregate):
|
||||
"""
|
||||
correct_prediction: True if the net_output (src_tokens) should
|
||||
predict the correct target
|
||||
aggregate: True if the criterion expects net_output (src_tokens)
|
||||
aggregated across time axis
|
||||
"""
|
||||
predicted_idx = 0 if correct_prediction else 1
|
||||
if aggregate:
|
||||
src_tokens = torch.zeros((2, 2), dtype=torch.float)
|
||||
for b in range(2):
|
||||
src_tokens[b][predicted_idx] = 1.0
|
||||
else:
|
||||
src_tokens = torch.zeros((2, 10, 2), dtype=torch.float)
|
||||
for b in range(2):
|
||||
for t in range(10):
|
||||
src_tokens[b][t][predicted_idx] = 1.0
|
||||
return src_tokens
|
||||
|
||||
def get_target(self, soft_target):
|
||||
if soft_target:
|
||||
target = torch.zeros((2, 2), dtype=torch.float)
|
||||
for b in range(2):
|
||||
target[b][0] = 1.0
|
||||
else:
|
||||
target = torch.zeros((2, 10), dtype=torch.long)
|
||||
return target
|
||||
|
||||
def get_test_sample(self, correct, soft_target, aggregate):
|
||||
src_tokens = self.get_src_tokens(correct, aggregate)
|
||||
target = self.get_target(soft_target)
|
||||
L = src_tokens.size(1)
|
||||
return {
|
||||
"net_input": {"src_tokens": src_tokens, "src_lengths": torch.tensor([L])},
|
||||
"target": target,
|
||||
"ntokens": src_tokens.size(0) * src_tokens.size(1),
|
||||
}
|
||||
Reference in New Issue
Block a user