Added Exportable bits to T5 + unit test

Signed-off-by: Boris Fomitchev <bfomitchev@nvidia.com>

nemo/collections/tts/g2p/models/t5.py

@@ -24,6 +24,7 @@
 from nemo.collections.asr.metrics.wer import word_error_rate
 from nemo.collections.tts.g2p.data.t5 import T5G2PDataset
 from nemo.collections.tts.models.base import G2PModel
+from nemo.core.classes import Exportable
 from nemo.core.classes.common import PretrainedModelInfo, typecheck
 from nemo.core.neural_types import LabelsType, LossType, MaskType, NeuralType, TokenIndex
 from nemo.utils import logging
@@ -38,7 +39,7 @@ class T5G2PConfig:
  test_ds: Optional[Dict[Any, Any]] = None
 
 
-class T5G2PModel(G2PModel):
+class T5G2PModel(G2PModel, Exportable):
  """
  T5-based grapheme-to-phoneme model.
  """
@@ -260,3 +261,12 @@ def setup_multiple_test_data(self, test_data_config: Union[DictConfig, Dict] = N
  @classmethod
  def list_available_models(cls) -> 'List[PretrainedModelInfo]':
  return []
+
+ # ===== export methods ===========$
+
+ def input_example(self, max_batch=1, max_dim=64, seq_len=16):
+ sample = next(self.parameters())
+ input_ids = torch.randint(low=0, high=max_dim, size=(max_batch, seq_len), device=sample.device)
+ labels = torch.randint(low=0, high=max_dim, size=(max_batch, seq_len), device=sample.device)
+ attention_mask = torch.randint(low=0, high=1, size=(max_batch, seq_len), device=sample.device)
+ return tuple([input_ids, attention_mask, labels])

nemo/core/classes/exportable.py

@@ -188,7 +188,9 @@ def _export(
  input_list, input_dict = parse_input_example(input_example)
  input_names = self.input_names
  output_names = self.output_names
- output_example = tuple(self.forward(*input_list, **input_dict))
+ output_example = self.forward(*input_list, **input_dict)
+ if torch.is_tensor(output_example):
+ output_example = (output_example,)
 
  if check_trace:
  if isinstance(check_trace, bool):

nemo/utils/export_utils.py

@@ -144,6 +144,8 @@ def verify_torchscript(model, output, input_examples, check_tolerance=0.01):
  # We disable autocast here to make sure exported TS will run under Triton or other C++ env
  with torch.cuda.amp.autocast(enabled=False):
  output_example = model.forward(*input_list, **input_dict)
+ if torch.is_tensor(output_example):
+ output_example = (output_example,)
  ts_model = torch.jit.load(output)
  all_good = all_good and run_ts_and_compare(
  ts_model, input_list, input_dict, output_example, check_tolerance
@@ -172,6 +174,8 @@ def verify_runtime(model, output, input_examples, input_names, check_tolerance=0
  for input_example in input_examples:
  input_list, input_dict = parse_input_example(input_example)
  output_example = model.forward(*input_list, **input_dict)
+ if torch.is_tensor(output_example):
+ output_example = (output_example,)
  ort_input = to_onnxrt_input(ort_input_names, input_names, input_dict, input_list)
  all_good = all_good and run_ort_and_compare(sess, ort_input, output_example, check_tolerance)
  status = "SUCCESS" if all_good else "FAIL"
@@ -184,6 +188,8 @@ def run_ts_and_compare(ts_model, ts_input_list, ts_input_dict, output_example, c
  ts_out = ts_model(*ts_input_list, **ts_input_dict)
 
  all_good = True
+ if torch.is_tensor(ts_out):
+ ts_out = (ts_out,)
  for i, out in enumerate(ts_out):
  expected = output_example[i]
 
@@ -206,6 +212,8 @@ def run_ort_and_compare(sess, ort_input, output_example, check_tolerance=0.01):
  # Verify the model can be read, and is valid
  ort_out = sess.run(None, ort_input)
  all_good = True
+ if torch.is_tensor(ort_out):
+ ort_out = (ort_out,)
  for i, out in enumerate(ort_out):
  expected = output_example[i]
 

tests/collections/tts/g2p/test_export.py

@@ -0,0 +1,65 @@
+# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import os
+import tempfile
+
+import pytest
+import torch
+from omegaconf import OmegaConf
+
+from nemo.collections.tts.g2p.models.t5 import T5G2PModel
+from nemo.utils.app_state import AppState
+
+
+@pytest.fixture()
+def t5_model():
+ this_test_dir = os.path.dirname(os.path.abspath(__file__))
+
+ cfg = OmegaConf.load(os.path.join(this_test_dir, '../../../../examples/tts/g2p/conf/g2p_t5.yaml'))
+ cfg.train_manifest = None
+ cfg.validation_manifest = None
+ app_state = AppState()
+ app_state.is_model_being_restored = True
+ model = T5G2PModel(cfg=cfg.model)
+ app_state.is_model_being_restored = False
+ model.eval()
+ return model
+
+
+def extra_cfg():
+ cfg.model.init_from_ptl_ckpt = None
+ cfg.model.train_ds.dataset.manifest_filepath = "dummy.json"
+ cfg.model.train_ds.dataset.sup_data_path = "dummy.json"
+ cfg.model.validation_ds.dataset.manifest_filepath = "dummy.json"
+ cfg.model.validation_ds.dataset.sup_data_path = "dummy.json"
+ cfg.pitch_mean = 212.35
+ cfg.pitch_std = 68.52
+
+
+class TestExportable:
+ @pytest.mark.run_only_on('GPU')
+ @pytest.mark.unit
+ def test_T5Model_export_to_onnx(self, t5_model):
+ model = t5_model.cuda()
+ with tempfile.TemporaryDirectory() as tmpdir:
+ filename = os.path.join(tmpdir, 'fp.onnx')
+ model.export(output=filename, verbose=True, onnx_opset_version=18, check_trace=True)
+
+ @pytest.mark.run_only_on('GPU')
+ @pytest.mark.unit
+ def test_T5Model_export_to_ts(self, t5_model):
+ model = t5_model.cuda()
+ with tempfile.TemporaryDirectory() as tmpdir:
+ filename = os.path.join(tmpdir, 'fp.ts')
+ model.export(output=filename, verbose=True, check_trace=True)