2020-01-22 05:24:00 +03:00
|
|
|
#!/usr/bin/env python3
|
2020-01-31 02:57:27 +03:00
|
|
|
# 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.
|
2020-01-22 05:24:00 +03:00
|
|
|
|
2020-01-31 02:57:27 +03:00
|
|
|
import argparse
|
2020-01-31 03:13:15 +03:00
|
|
|
import tempfile
|
2020-01-22 05:24:00 +03:00
|
|
|
import unittest
|
|
|
|
|
|
|
|
import torch
|
2021-11-29 23:30:10 +03:00
|
|
|
|
2020-01-31 02:57:27 +03:00
|
|
|
from fairseq.data.dictionary import Dictionary
|
|
|
|
from fairseq.models.transformer import TransformerModel
|
2020-01-22 21:53:43 +03:00
|
|
|
from fairseq.modules import multihead_attention, sinusoidal_positional_embedding
|
2020-09-10 03:00:56 +03:00
|
|
|
from fairseq.tasks.fairseq_task import LegacyFairseqTask
|
2020-01-22 05:24:00 +03:00
|
|
|
|
2020-01-31 02:57:27 +03:00
|
|
|
DEFAULT_TEST_VOCAB_SIZE = 100
|
|
|
|
|
|
|
|
|
2020-09-10 03:00:56 +03:00
|
|
|
class DummyTask(LegacyFairseqTask):
|
2020-01-31 02:57:27 +03:00
|
|
|
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.src_dict = self.dictionary
|
|
|
|
self.tgt_dict = self.dictionary
|
|
|
|
|
|
|
|
@property
|
|
|
|
def source_dictionary(self):
|
|
|
|
return self.src_dict
|
|
|
|
|
|
|
|
@property
|
|
|
|
def target_dictionary(self):
|
|
|
|
return self.dictionary
|
|
|
|
|
|
|
|
|
|
|
|
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
|
2020-01-27 21:23:56 +03:00
|
|
|
|
2020-01-31 02:57:27 +03:00
|
|
|
|
|
|
|
def get_dummy_task_and_parser():
|
|
|
|
"""
|
|
|
|
Return a dummy task and argument parser, which can be used to
|
|
|
|
create a model/criterion.
|
|
|
|
"""
|
|
|
|
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
|
|
|
|
|
|
|
|
|
2020-04-11 04:17:09 +03:00
|
|
|
def _test_save_and_load(scripted_module):
|
|
|
|
with tempfile.NamedTemporaryFile() as f:
|
|
|
|
scripted_module.save(f.name)
|
|
|
|
torch.jit.load(f.name)
|
|
|
|
|
2020-01-31 03:13:15 +03:00
|
|
|
|
2020-04-11 04:17:09 +03:00
|
|
|
class TestExportModels(unittest.TestCase):
|
2020-01-22 05:24:00 +03:00
|
|
|
def test_export_multihead_attention(self):
|
|
|
|
module = multihead_attention.MultiheadAttention(embed_dim=8, num_heads=2)
|
2020-01-31 03:13:15 +03:00
|
|
|
scripted = torch.jit.script(module)
|
2020-04-11 04:17:09 +03:00
|
|
|
_test_save_and_load(scripted)
|
2020-01-22 21:53:43 +03:00
|
|
|
|
2020-01-27 21:23:56 +03:00
|
|
|
def test_incremental_state_multihead_attention(self):
|
|
|
|
module1 = multihead_attention.MultiheadAttention(embed_dim=8, num_heads=2)
|
|
|
|
module1 = torch.jit.script(module1)
|
|
|
|
module2 = multihead_attention.MultiheadAttention(embed_dim=8, num_heads=2)
|
|
|
|
module2 = torch.jit.script(module2)
|
|
|
|
|
|
|
|
state = {}
|
2020-01-31 02:57:27 +03:00
|
|
|
state = module1.set_incremental_state(state, "key", {"a": torch.tensor([1])})
|
|
|
|
state = module2.set_incremental_state(state, "key", {"a": torch.tensor([2])})
|
|
|
|
v1 = module1.get_incremental_state(state, "key")["a"]
|
|
|
|
v2 = module2.get_incremental_state(state, "key")["a"]
|
2020-01-27 21:23:56 +03:00
|
|
|
|
|
|
|
self.assertEqual(v1, 1)
|
|
|
|
self.assertEqual(v2, 2)
|
|
|
|
|
2020-01-22 21:53:43 +03:00
|
|
|
def test_positional_embedding(self):
|
|
|
|
module = sinusoidal_positional_embedding.SinusoidalPositionalEmbedding(
|
|
|
|
embedding_dim=8, padding_idx=1
|
|
|
|
)
|
2020-01-31 03:13:15 +03:00
|
|
|
scripted = torch.jit.script(module)
|
2020-04-11 04:17:09 +03:00
|
|
|
_test_save_and_load(scripted)
|
2020-01-31 02:57:27 +03:00
|
|
|
|
|
|
|
@unittest.skipIf(
|
2022-08-02 00:26:36 +03:00
|
|
|
torch.__version__ < "1.6.0", "Targeting OSS scriptability for the 1.6 release"
|
2020-01-31 02:57:27 +03:00
|
|
|
)
|
|
|
|
def test_export_transformer(self):
|
|
|
|
task, parser = get_dummy_task_and_parser()
|
|
|
|
TransformerModel.add_args(parser)
|
|
|
|
args = parser.parse_args([])
|
|
|
|
model = TransformerModel.build_model(args, task)
|
2020-01-31 03:13:15 +03:00
|
|
|
scripted = torch.jit.script(model)
|
2020-04-11 04:17:09 +03:00
|
|
|
_test_save_and_load(scripted)
|
2020-01-31 02:57:27 +03:00
|
|
|
|
2021-02-23 01:21:36 +03:00
|
|
|
@unittest.skipIf(
|
2022-08-02 00:26:36 +03:00
|
|
|
torch.__version__ < "1.6.0", "Targeting OSS scriptability for the 1.6 release"
|
2021-02-23 01:21:36 +03:00
|
|
|
)
|
|
|
|
def test_export_transformer_no_token_pos_emb(self):
|
|
|
|
task, parser = get_dummy_task_and_parser()
|
|
|
|
TransformerModel.add_args(parser)
|
|
|
|
args = parser.parse_args([])
|
|
|
|
args.no_token_positional_embeddings = True
|
|
|
|
model = TransformerModel.build_model(args, task)
|
|
|
|
scripted = torch.jit.script(model)
|
|
|
|
_test_save_and_load(scripted)
|
|
|
|
|
|
|
|
|
2020-01-31 02:57:27 +03:00
|
|
|
if __name__ == "__main__":
|
|
|
|
unittest.main()
|