Hello everyone, not long ago we launched the model deployment tutorial series, which has been praised by everyone, and also received a lot of friends’ urging. The follow-up tutorial is being prepared, and we will meet with you soon, please look forward to it

OpenMMLab: Model Deployment Tutorial (part I) : Introduction to Model Deployment

OpenMMLab: Getting started with Model Deployment (Part 2) : Solving problems in Model Deployment

Today, we kick off a new TorchScript interpretation series that takes you through the PyTorch model deployment.

What is TorchScript?

PyTorch is undoubtedly one of the most successful training frameworks for deep learning today, and is a hot topic for all kinds of top paper experiments. Compared to other frameworks, PyTorch’s biggest selling point is its support for dynamic networks, which has a lower learning cost than other frameworks that require static networks to be built. The PyTorch source Readme also includes a dynamic diagram for this purpose:

PyTorch is a great tool for researchers to come up with ideas, do experiments, and write papers, but it is not suitable for deployment. The advantage of dynamic drawing is more like a disadvantage for more performance-demanding applications. The unfixed network structure makes it difficult to analyze and optimize the network structure, most parameters can be transmitted in Tensor form and resource allocation becomes a headache. In addition, since diagrams are built from Python code, deployment is dependent on the Python environment, and there is no confidentiality in the model.

TorchScript was created to solve this problem. Including code tracing and parsing, intermediate representation generation, model optimization, serialization and other functions, can be said to cover all aspects of model deployment. Today we will briefly introduce some of the features of TorchScript, so that you can have a preliminary understanding of the advanced interpretation will be introduced

Model transformation

As a paradigm for model deployment, we usually need to generate an intermediate representation (IR) of the model, which has a relatively fixed graph structure and is therefore easier to optimize. Let’s look at an example:

Model = resnet18() model.eval() import torch from Torchvision. models import resnet18 # Dummy_input = torch. Rand (1, 3, 224, 224); jit_model = torch.jit.trace(model, dummy_input)Copy the code

At this point, PyTorch’s model is converted to TorchScript’s IR. Here, we use the trace mode to generate IR. The so-called trace means to conduct a model inference, record all the calculations during the inference process, and integrate these records into a calculation graph. We will explain the process of trace in future sharing.

So what does this IR have in it? We can visualize layer1:

jit_layer1 = jit_model.layer1 
print(jit_layer1.graph) 
 
# graph(%self.6 : __torch__.torch.nn.modules.container.Sequential, 
#       %4 : Float(1, 64, 56, 56, strides=[200704, 3136, 56, 1], requires_grad=0, device=cpu)): 
#   %1 : __torch__.torchvision.models.resnet.___torch_mangle_10.BasicBlock = prim::GetAttr[name="1"](%self.6) 
#   %2 : __torch__.torchvision.models.resnet.BasicBlock = prim::GetAttr[name="0"](%self.6) 
#   %6 : Tensor = prim::CallMethod[name="forward"](%2, %4) 
#   %7 : Tensor = prim::CallMethod[name="forward"](%1, %6) 
#   return (%7) 
Copy the code

A little confused? TorchScript has its own definition of Graph and its elements, which may be unfamiliar to those who are new to it, but that’s ok, there’s another way to visualize it:

print(jit_layer1.code) 
 
# def forward(self, 
#     argument_1: Tensor) -> Tensor: 
#   _0 = getattr(self, "1") 
#   _1 = (getattr(self, "0")).forward(argument_1, ) 
#   return (_0).forward(_1, ) 
Copy the code

Yes, code! TorchScript IR can be restored to Python code, so if you generate a TorchScript model and want to know if it has the right contents, you can do some simple checks like this.

So in the example we just used trace to generate IR. In addition to trace, PyTorch provides another way to generate TorchScript models: script. This approach directly parses the Python code of the network definition and generates an abstract syntax tree AST. Therefore, this approach can solve some problems that cannot be solved by trace, such as mapping data flow control statements such as Branch /loop. Script mode of drawing has many interesting features, will do a special analysis in the future sharing, please look forward to.

Model optimization

If you are smart enough, you may have noticed that only the forward part of resnet18 is in the visualization above. Is the submodule information missing? If not, how do you determine if the submodule’s contents are correct? Don’t worry, remember we said that TorchScript supports network optimization, here we can solve this problem with a pass:

Call inline pass, 12 torch. _c._jit_pass_inline (jit_layer1.graph) print(jit_layer1.code) # def forward(self, # argument_1: Tensor) -> Tensor: # _0 = getattr(self, "1") # _1 = getattr(self, "0") # _2 = _1.bn2 # _3 = _1.conv2 # _4 = _1.bn1 # input = torch._convolution(argument_1, _1.conv1.weight, None, [1, 1], [1, 1], [1, 1], False, [0, 0], 1, False, False, True, True) # _5 = _4.running_var # _6 = _4.running_mean # _7 = _4.bias # input0 = torch.batch_norm(input, _4.weight, _7, _6, _5, False, 0.10000000000000001, 1.0000000000000001 e-05, True) # input1 = torch.relu_(input0) # input2 = torch._convolution(input1, _3.weight, None, [1, 1], [1, 1], [1, 1], False, [0, 0], 1, False, False, True, True) # _8 = _2.running_var # _9 = _2.running_mean # _10 = _2.bias # out = torch.batch_norm(input2, _2.weight, _10, _9, _8, False, 0.100000000000000000001, 1.0000000000000001E-05, True) # input3 = torch. Add_ (out, Argument_1, alpha=1) # input4 = torch.relu_(input3) # _11 = _0.bn2 # _12 = _0.conv2 # _13 = _0.bn1 # input5 = torch._convolution(input4, _0.conv1.weight, None, [1, 1], [1, 1], [1, 1], False, [0, 0], 1, False, False, True, True) # _14 = _13.running_var # _15 = _13.running_mean # _16 = _13.bias # input6 = torch.batch_norm(input5, _13.weight, _16, _15, _14, False, 0.100000000000000000001, 1.0000000000000001E-05, True) # input7 = torch.relu_(input6) # input8 = torch._convolution(input7, _12.weight, None, [1, 1], [1, 1], [1, 1], False, [0, 0], 1, False, False, True, True) # _17 = _11.running_var # _18 = _11.running_mean # _19 = _11.bias # out0 = torch.batch_norm(input8, _11.weight, _19, _18, _17, False, 0.10000000000000001, 1.000000000001E-05, True) # input9 = torch. Add_ (out0, input4, alpha=1) # return torch.relu_(input9)Copy the code

Here we see familiar operators like convolution, batch_norm, relu, etc.

In the code above we use a pass named inline to inline all submodules so we can see the more complete reasoning code. Pass is a concept derived from compilation principles. A TorchScript pass takes a graph, iterates over all the elements in the graph, performs some kind of transformation, and generates a new graph. The inline we use here expands the module call, and although this does not directly affect execution efficiency, it is the basis for many other passes. There are a number of passes defined in PyTorch to solve various optimization tasks, which we’ll cover in more detail in the future.

serialization

TorchScript can be serialized either way, for example:

Jit_model = torch. Jit. Load ('jit_model.Copy the code

The serialized model is no longer python-specific and can be deployed to a variety of platforms.

PyTorch provides a c++ API for TorchScript model inference, so serialized models can finally be inferred without relying on python:

// Load the generated Torchscript model auto Module = Torch ::jit::load('jit_model.pth'); STD ::vector< Torch :: JIT ::IValue> inputs = inputs STD ::vector<torch:: JIT ::IValue> inputs = inputs ; Auto output = module. Forward (inputs). ToTensor ();Copy the code

Relationships with other components?

Relationship with Torch. Onnx?

ONNX is a neural network intermediate representation widely used in the industry, and PyTorch naturally supports ONNX as well. The torch. Onnx. export function helps convert the PyTorch model to an OnNx model by tracing PyTorch’s reasoning. As you might already know, yes, the ONNX export uses TorchScript trace. The specific steps are as follows:

  1. Use trace to create a TorchScipt model, you can skip this step if you are converting TorchScript.
  2. Transform the model generated in 1 using a number of passes, of which the most important pass for ONNX exports isToONNXThat pass will do a mapping to TorchScriptprim,atenThe operator in space maps toonnxOperator in space.
  3. The model is serialized using THE PROto format of ONNX to complete the export of ONNX.

The implementation of ONNX export and operator mapping will be discussed in more detail in future shares.

The relationship with Torch. Fx?

PyTorch1.9 began with the addition of the torch. Fx tool, which, according to the official description, consists of symbolic tracer, intermediate representation, Python Code Generation, and other components. Python -> Python translation implemented. Does it look a little bit like TorchScript?

In fact, they are not closely related to each other. They can be regarded as two perpendicular tools, created to solve two different tasks.

  • The main use of TorchScript is for model deployment, which requires recording to generate an IR for inference optimization, editing of computational diagrams is often performance-oriented, and so on, without adding new functionality to the model itself.
  • The main use of FX is to conductpython->pythonIts IR node types are simpler, for exampleA function call,The extractionWait, this kind of IR learning is cheaper and easier to edit. Using FX to edit diagrams is usually for some specific purpose, such as inserting quantization nodes into the model, to avoid the rework of manually editing the network.

The two tools can be used together, such as using the FX tool to edit models to make training easier and more powerful. TorchScript is then used to accelerate the deployment of the model to a specific platform.

I hope that through the above sharing, you have a preliminary understanding of TorchScript, we will bring you more advanced interpretation in the future, welcome to continue to pay attention to. TorchScript is now supported by MMDeploy. You are welcome to visit MMDeploy’s GitHub home page to experience it.

This is the end of today’s sharing, have you learned? If we share to bring you certain help, welcome to like collect attention, than heart ~