JIT in WeNet

We want that our PyTorch model can be directly exported by torch.jit.script method, which is essential for deploying the model to production.

See the following resource for how to deploy PyTorch models in production in details.

To ensure that, we will try to export the model before training stage. If it fails, we should modify the training code to satisfy the export requirements.

# See in wenet/bin/train.py
script_model = torch.jit.script(model)
script_model.save(os.path.join(args.model_dir, 'init.zip'))

Two principles should be taken into consideration when we contribute our python code to WeNet, especially for the subclass of torch.nn.Module, and for the forward function.

  1. Know what is allowed and what is disallowed.

  2. Try to use explicit typing as much as possible. You can try to do type checking forced by typeguard, see https://typeguard.readthedocs.io/en/latest/userguide.html for details.