Add save and load tests to fairseq export test
Summary:
Pull Request resolved:
Earlier we had some issues at pickling. Type information gets lost. Fixed in
These save_and_load tests are added for protection in the future.
# LICENSE file in the root directory of this source tree.
import argparse
+import tempfile
import unittest
import torch
@@ -58,9 +59,15 @@ def get_dummy_task_and_parser():
class TestExportModels(unittest.TestCase):
+ def _test_save_and_load(self, scripted_module):
+ with tempfile.NamedTemporaryFile() as f:
+ scripted_module.save(f.name)
+ torch.jit.load(f.name)
+
def test_export_multihead_attention(self):
module = multihead_attention.MultiheadAttention(embed_dim=8, num_heads=2)
- torch.jit.script(module)
+ scripted = torch.jit.script(module)
+ self._test_save_and_load(scripted)
def test_incremental_state_multihead_attention(self):
module1 = multihead_attention.MultiheadAttention(embed_dim=8, num_heads=2)
@@ -81,7 +88,8 @@ class TestExportModels(unittest.TestCase):
module = sinusoidal_positional_embedding.SinusoidalPositionalEmbedding(
embedding_dim=8, padding_idx=1
)
- torch.jit.script(module)
+ scripted = torch.jit.script(module)
+ self._test_save_and_load(scripted)
@unittest.skipIf(
torch.__version__ < "1.5.0", "Targeting OSS scriptability for the 1.5 release"
@@ -91,7 +99,8 @@ class TestExportModels(unittest.TestCase):
TransformerModel.add_args(parser)
args = parser.parse_args([])
model = TransformerModel.build_model(args, task)
- torch.jit.script(model)
+ scripted = torch.jit.script(model)
+ self._test_save_and_load(scripted)
if __name__ == "__main__":
