Logo Questions Linux Laravel Mysql Ubuntu Git Menu
 

What are Torch Scripts in PyTorch?

Tags:

jit

pytorch

I've just found that PyTorch docs expose something that is called Torch Scripts. However, I do not know:

  • When they should be used?
  • How they should be used?
  • What are their benefits?
like image 277
ndrwnaguib Avatar asked Dec 23 '18 00:12

ndrwnaguib


People also ask

What does torch JIT do?

PyTorch JIT is an optimizing JIT compiler for PyTorch. It uses runtime information to optimize TorchScript modules. It can automate optimizations like layer fusion, quantization, sparsification.

How do you convert PyTorch to TorchScript?

There are two ways to convert your model to TorchScript: tracing and scripting. We will only demonstrate the first one, tracing, but you can find information about scripting from the PyTorch documentation. When tracing, we use an example input to record the actions taken and capture the the model architecture.

Can you train a TorchScript model?

Torchscript main purpose is to run models in production environments in inference mode. It is not designed for training networks, you should use the Pytorch code you used for training instead.


1 Answers

Torch Script is one of two modes of using the PyTorch just in time compiler, the other being tracing. The benefits are explained in the linked documentation:

Torch Script is a way to create serializable and optimizable models from PyTorch code. Any code written in Torch Script can be saved from your Python process and loaded in a process where there is no Python dependency.

The above quote is actually true both of scripting and tracing. So

  1. You gain the ability to serialize your models and later run them outside of Python, via LibTorch, a C++ native module. This allows you to embed your DL models in various production environments like mobile or IoT. There is an official guide on exporting models to C++ here.
  2. PyTorch can compile your jit-able modules rather than running them as an interpreter, allowing for various optimizations and improving performance, both during training and inference. This is equally helpful for development and production.

Regarding Torch Script specifically, in comparison to tracing, it is a subset of Python, specified in detail here, which, when adhered to, can be compiled by PyTorch. It is more laborious to write Torch Script modules instead of tracing regular nn.Module subclasses, but it allows for some extra features over tracing, most notably flow control like if statements or for loops. Tracing treats such flow control as "constant" - in other words, if you have an if model.training clause in your module and trace it with training=True, it will always behave this way, even if you change the training variable to False later on.

To answer your first question, you need to use jit if you want to deploy your models outside Python and otherwise you should use jit if you want to gain some execution performance at the price of extra development effort (as not every model can be straightforwardly made compliant with jit). In particular, you should use Torch Script if your code cannot be jited with tracing alone because it relies on some features such as if statements. For maximum ergonomy, you probably want to mix the two on a case-by-case basis.

Finally, for how they should be used, please refer to all the documentation and tutorial links.

like image 174
Jatentaki Avatar answered Sep 19 '22 06:09

Jatentaki