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.