Pytorch1.0 use torch script to export and save the Model

Python's ease of use and the dynamic graph feature of pytorch make pytorch more and more popular in academic research. However, in the production environment, due to Python's Gil and other features, high concurrency and low latency requirements may not be met, and a C ++ interface is required. In addition to exporting the model to onnx, pytorch1.0 provides a new solution: pytorch Training Model-saving the model through the torch script intermediate script-C ++ Loading Model. In recent work, we need to try to make a conversion and summarize the steps and pitfalls.

There are two ways to convert the torch model into a C ++ interface readable model using torch Script: trace & script. trace is simpler than script, but it is only suitable for fixed network models, that is, there is no control flow in forward, because trace only saves the actual path when running. If the forward function has a control flow, it must be implemented in script mode.

As the name suggests, trace is to go through the data operation path. Official example:

import torchdef foo(x, y): return 2*x + ytraced_foo = torch.jit.trace(foo, (torch.rand(3), torch.rand(3)))

 

 

 

 

The script is a little complicated and mainly involves three changes:

1. The model inherits from Nn. Model to torch. JIT. scriptmodule.

2. Add @ torch. JIT. script_method before the forward Function

3. Add @ torch. JIT. script before other functions to be called.

 

Pitfalls and solutions:

A. Torch script default function or method parameters are of the tensor type. If this parameter is not required, a compilation error of Type Mismatch will be reported when a non-tensor parameter is called.

Python3 can be directly:

Def example_func (param_1: tensor, param_2: int, param_3: list [int]):

 

 

Python2 needs to be annotated with type:

Def example_func (param_1, param_2, param_3 ): # Type: (tensor, Int, list [int])-> Tensor

 

 

 

BIn the. model method, orward adds @ torch. JIT. script_method, And the _ init _ function does not need to be used.

C. As mentioned above, torch scrip supports a subset of pytorch, which means some functions are not supported, such as not Boolean, pass, list slice assignment, and CPU and GPU switching value. to (), you need to find a way to bypass. In the discussion board on GitHub, it seems that the new version has supported the not operation and has not been verified.

 

Conclusion: The current preview version of pytorch 1.0 has a lot of room for optimization, at least in the function set supported by torch script. It is not recommended to use it. Let's wait for the stable version to be released.

　　

 

Original content. For more information, see the source.

 

References:

Https://pytorch.org/docs/master/jit.html

Https://pytorch.org/tutorials/beginner/deploy_seq2seq_hybrid_frontend_tutorial.html

Pytorch1.0 use torch script to export and save the Model