Add tests to ensure that the model is torch scriptable.

egs/librispeech/ASR/pruned_transducer_stateless/test_model.py

@@ -0,0 +1,50 @@
+#!/usr/bin/env python3
+# Copyright    2022  Xiaomi Corp.        (authors: Fangjun Kuang)
+#
+# See ../../../../LICENSE for clarification regarding multiple authors
+#
+# 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.
+
+
+"""
+To run this file, do:
+
+    cd icefall/egs/librispeech/ASR
+    python ./pruned_transducer_stateless/test_model.py
+"""
+
+import torch
+from train import get_params, get_transducer_model
+
+
+def test_model():
+    params = get_params()
+    params.vocab_size = 500
+    params.blank_id = 0
+    params.context_size = 2
+    params.unk_id = 2
+
+    model = get_transducer_model(params)
+
+    num_param = sum([p.numel() for p in model.parameters()])
+    print(f"Number of model parameters: {num_param}")
+    model.__class__.forward = torch.jit.ignore(model.__class__.forward)
+    torch.jit.script(model)
+
+
+def main():
+    test_model()
+
+
+if __name__ == "__main__":
+    main()

egs/librispeech/ASR/pruned_transducer_stateless2/test_model.py

@@ -0,0 +1,50 @@
+#!/usr/bin/env python3
+# Copyright    2022  Xiaomi Corp.        (authors: Fangjun Kuang)
+#
+# See ../../../../LICENSE for clarification regarding multiple authors
+#
+# 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.
+
+
+"""
+To run this file, do:
+
+    cd icefall/egs/librispeech/ASR
+    python ./pruned_transducer_stateless2/test_model.py
+"""
+
+import torch
+from train import get_params, get_transducer_model
+
+
+def test_model():
+    params = get_params()
+    params.vocab_size = 500
+    params.blank_id = 0
+    params.context_size = 2
+    params.unk_id = 2
+
+    model = get_transducer_model(params)
+
+    num_param = sum([p.numel() for p in model.parameters()])
+    print(f"Number of model parameters: {num_param}")
+    model.__class__.forward = torch.jit.ignore(model.__class__.forward)
+    torch.jit.script(model)
+
+
+def main():
+    test_model()
+
+
+if __name__ == "__main__":
+    main()

egs/librispeech/ASR/pruned_transducer_stateless3/test_model.py

@@ -0,0 +1,50 @@
+#!/usr/bin/env python3
+# Copyright    2022  Xiaomi Corp.        (authors: Fangjun Kuang)
+#
+# See ../../../../LICENSE for clarification regarding multiple authors
+#
+# 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.
+
+
+"""
+To run this file, do:
+
+    cd icefall/egs/librispeech/ASR
+    python ./pruned_transducer_stateless3/test_model.py
+"""
+
+import torch
+from train import get_params, get_transducer_model
+
+
+def test_model():
+    params = get_params()
+    params.vocab_size = 500
+    params.blank_id = 0
+    params.context_size = 2
+    params.unk_id = 2
+
+    model = get_transducer_model(params)
+
+    num_param = sum([p.numel() for p in model.parameters()])
+    print(f"Number of model parameters: {num_param}")
+    model.__class__.forward = torch.jit.ignore(model.__class__.forward)
+    torch.jit.script(model)
+
+
+def main():
+    test_model()
+
+
+if __name__ == "__main__":
+    main()

egs/librispeech/ASR/pruned_transducer_stateless4/test_model.py

@@ -0,0 +1,50 @@
+#!/usr/bin/env python3
+# Copyright    2022  Xiaomi Corp.        (authors: Fangjun Kuang)
+#
+# See ../../../../LICENSE for clarification regarding multiple authors
+#
+# 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.
+
+
+"""
+To run this file, do:
+
+    cd icefall/egs/librispeech/ASR
+    python ./pruned_transducer_stateless4/test_model.py
+"""
+
+import torch
+from train import get_params, get_transducer_model
+
+
+def test_model():
+    params = get_params()
+    params.vocab_size = 500
+    params.blank_id = 0
+    params.context_size = 2
+    params.unk_id = 2
+
+    model = get_transducer_model(params)
+
+    num_param = sum([p.numel() for p in model.parameters()])
+    print(f"Number of model parameters: {num_param}")
+    model.__class__.forward = torch.jit.ignore(model.__class__.forward)
+    torch.jit.script(model)
+
+
+def main():
+    test_model()
+
+
+if __name__ == "__main__":
+    main()

egs/librispeech/ASR/transducer_stateless/joiner.py

@@ -14,6 +14,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from typing import List
+
 import torch
 import torch.nn as nn
 
@@ -55,8 +57,8 @@ class Joiner(nn.Module):
 
         N = encoder_out.size(0)
 
-        encoder_out_len = encoder_out_len.tolist()
+        encoder_out_len: List[int] = encoder_out_len.tolist()
-        decoder_out_len = decoder_out_len.tolist()
+        decoder_out_len: List[int] = decoder_out_len.tolist()
 
         encoder_out_list = [
             encoder_out[i, : encoder_out_len[i], :] for i in range(N)

egs/librispeech/ASR/transducer_stateless/test_model.py

@@ -0,0 +1,49 @@
+#!/usr/bin/env python3
+# Copyright    2022  Xiaomi Corp.        (authors: Fangjun Kuang)
+#
+# See ../../../../LICENSE for clarification regarding multiple authors
+#
+# 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.
+
+
+"""
+To run this file, do:
+
+    cd icefall/egs/librispeech/ASR
+    python ./transducer_stateless/test_model.py
+"""
+
+import torch
+from train import get_params, get_transducer_model
+
+
+def test_model():
+    params = get_params()
+    params.vocab_size = 500
+    params.blank_id = 0
+    params.context_size = 2
+
+    model = get_transducer_model(params)
+
+    num_param = sum([p.numel() for p in model.parameters()])
+    print(f"Number of model parameters: {num_param}")
+    model.__class__.forward = torch.jit.ignore(model.__class__.forward)
+    torch.jit.script(model)
+
+
+def main():
+    test_model()
+
+
+if __name__ == "__main__":
+    main()

egs/librispeech/ASR/transducer_stateless2/joiner.py

@@ -14,6 +14,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from typing import Optional
+
 import torch
 import torch.nn as nn
 
@@ -30,7 +32,8 @@ class Joiner(nn.Module):
         self,
         encoder_out: torch.Tensor,
         decoder_out: torch.Tensor,
-        *unused,
+        unused_encoder_out_len: Optional[torch.Tensor] = None,
+        unused_decoder_out_len: Optional[torch.Tensor] = None,
     ) -> torch.Tensor:
         """
         Args:
@@ -38,10 +41,12 @@ class Joiner(nn.Module):
             Output from the encoder. Its shape is (N, T, self.input_dim).
           decoder_out:
             Output from the decoder. Its shape is (N, U, self.input_dim).
-          unused:
+          unused_encoder_out_len:
             This is a placeholder so that we can reuse
             transducer_stateless/beam_search.py in this folder as that
             script assumes the joiner networks accepts 4 inputs.
+          unused_decoder_out_len:
+            Just a placeholder.
         Returns:
           Return a tensor of shape (N, T, U, self.output_dim).
         """

egs/librispeech/ASR/transducer_stateless2/test_model.py

@@ -0,0 +1,49 @@
+#!/usr/bin/env python3
+# Copyright    2022  Xiaomi Corp.        (authors: Fangjun Kuang)
+#
+# See ../../../../LICENSE for clarification regarding multiple authors
+#
+# 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.
+
+
+"""
+To run this file, do:
+
+    cd icefall/egs/librispeech/ASR
+    python ./transducer_stateless2/test_model.py
+"""
+
+import torch
+from train import get_params, get_transducer_model
+
+
+def test_model():
+    params = get_params()
+    params.vocab_size = 500
+    params.blank_id = 0
+    params.context_size = 2
+
+    model = get_transducer_model(params)
+
+    num_param = sum([p.numel() for p in model.parameters()])
+    print(f"Number of model parameters: {num_param}")
+    model.__class__.forward = torch.jit.ignore(model.__class__.forward)
+    torch.jit.script(model)
+
+
+def main():
+    test_model()
+
+
+if __name__ == "__main__":
+    main()