LibTorch實戰五：模型序列化

一、在C++環境中加載一個TORCHSCRIP

官網：https://pytorch.org/tutorials/advanced/cpp_export.html

As its name suggests, the primary interface to PyTorch is the Python programming language. While Python is a suitable and preferred language for many scenarios requiring dynamism and ease of iteration, there are equally many situations where precisely these properties of Python are unfavorable. One environment in which the latter often applies is production – the land of low latencies and strict deployment requirements. For production scenarios, C++ is very often the language of choice, even if only to bind it into another language like Java, Rust or Go. The following paragraphs will outline the path PyTorch provides to go from an existing Python model to a serialized representation that can be loaded and executed purely from C++, with no dependency on Python.

　　一般地，類似python的腳本語言可用於算法快速實現、驗證；但在產品化過程中，一般采用效率更高的C++語言，下面的工作就是將模型從python環境中移植到c++環境。

 

Step1： 將pytorch模型轉為torch scrip類型的模型

A PyTorch model’s journey from Python to C++ is enabled by Torch Script, a representation of a PyTorch model that can be understood, compiled and serialized by the Torch Script compiler. If you are starting out from an existing PyTorch model written in the vanilla “eager” API, you must first convert your model to Torch Script. In the most common cases, discussed below, this requires only little effort. If you already have a Torch Script module, you can skip to the next section of this tutorial.

通過TorchSript，我們可將pytorch模型從python轉為c++。那么，什么是TorchScript呢？其實，它也是Pytorch模型的一種，這種模型能夠被TorchScript的編譯器識別讀取、序列化。一般地，在處理模型過程中，我們都會先將模型轉為torch script格式，例如：".pt"  -> "yolov5x.torchscript.pt"

There exist two ways of converting a PyTorch model to Torch Script. The first is known as tracing, a mechanism in which the structure of the model is captured by evaluating it once using example inputs, and recording the flow of those inputs through the model. This is suitable for models that make limited use of control flow. The second approach is to add explicit annotations to your model that inform the Torch Script compiler that it may directly parse and compile your model code, subject to the constraints imposed by the Torch Script language.

轉為torchscript格式有兩種方法：一是函數torch.jit.trace；二是函數torch.jit.script。

torch.jit.trace原理：基於跟蹤機制，需要輸入一張圖(0矩陣、張量亦可)，模型會對輸入的tensor進行處理，並記錄所有張量的操作，torch::jit::trace能夠捕獲模型的結構、參數並保存。由於跟蹤僅記錄張量上的操作，因此它不會記錄任何控制流操作，如if語句或循環。

torch.jit.script原理：需要開發者先定義好神經網絡模型結構，即：提前寫好 class MyModule(torch.nn.Module)，這樣TorchScript可以根據定義好的MyModule來解析網絡結構。

 

基於Tracing的方法來轉換為Torch Script

如下代碼，給 torch.jit.trace 函數輸入一個指定size的隨機張量、ResNet18的網絡模型，得到一個類型為 torch.jit.ScriptModule 的對象，即：traced_script_module

import torch
import torchvision

# An instance of your model.
model = torchvision.models.resnet18()

# An example input you would normally provide to your model's forward() method.
example = torch.rand(1, 3, 224, 224)

# Use torch.jit.trace to generate a torch.jit.ScriptModule via tracing.
traced_script_module = torch.jit.trace(model, example)

經過上述處理，traced_script_module變量已經包含網絡的結構和參數，可以直接用於推理，如下代碼：

In[1]: output = traced_script_module(torch.ones(1, 3, 224, 224))
In[2]: output[0, :5]
Out[2]: tensor([-0.2698, -0.0381,  0.4023, -0.3010, -0.0448], grad_fn=<SliceBackward>)

基於Annotating (Script)的方法來轉換為Torch Script

如果你的模型中有類似於控制流操作(例如：if or for循環)，基於上述tracing的方式不再適用，這種方式會排上用場，下面以vanilla模型為例子，注：下面網絡結構中有個if判斷。

# 定義一個vanilla模型
import torch

class MyModule(torch.nn.Module):
    def __init__(self, N, M):
        super(MyModule, self).__init__()
        self.weight = torch.nn.Parameter(torch.rand(N, M))

    def forward(self, input):
        if input.sum() > 0:
          output = self.weight.mv(input)
        else:
          output = self.weight + input
        return output

這里調用 torch.jit.script 來獲取 torch.jit.ScriptModule 類型的對象，即：sm

class MyModule(torch.nn.Module):
    def __init__(self, N, M):
        super(MyModule, self).__init__()
        self.weight = torch.nn.Parameter(torch.rand(N, M))

    def forward(self, input):
        if input.sum() > 0:
          output = self.weight.mv(input)
        else:
          output = self.weight + input
        return output

my_module = MyModule(10,20)
sm = torch.jit.script(my_module)

 

Step2： 序列化torch.jit.ScriptModule類型的對象，並保存為文件

注：上述的tacing和script方法都將得到一個類型為torch.jit.ScriptModule的對象(這里簡單記為：ScriptModule )，該對象就是常規的前向傳播模塊。不管是哪一種方法，此時，只需要將ScriptModule進行序列化保存就行。這里保存的是上述基於Tracing得到的ResNet推理模塊traced_script_module。

traced_script_module.save("traced_resnet_model.pt") # 序列化，保存
# 保存后可用工具：https://netron.app/ 進行可視化

同理，如下是保存基於Annotating得到推理模塊my_module 后續，在libtorch中加載上述保存的模型文件就行，不再依賴任何python包。

my_module.save("my_module_model.pt") # 為什么不是sm

 

Step3： 在libtorch中加載ScriptModule模型

 如何配置libtorh？，我這里僅貼下vs環境下的屬性表：

include：
D:\ThirdParty\libtorch-win-shared-with-deps-1.7.1+cu110\libtorch\include
D:\ThirdParty\libtorch-win-shared-with-deps-1.7.1+cu110\libtorch\include\torch\csrc\api\include

lib：
D:\ThirdParty\libtorch-win-shared-with-deps-1.7.1+cu110\libtorch\lib

鏈接器：
c10.lib
c10_cuda.lib
torch.lib
torch_cpu.lib
torch_cuda.lib

環境變量：
D:\ThirdParty\libtorch-win-shared-with-deps-1.7.1+cu110\libtorch\lib

以下c++代碼加載上述模型文件

#include<torch/torch.h>
#include<torch/script.h>
#include<iostream>
#include<memory>

int main()
{
    torch::jit::script::Module module;
    std::string str = "traced_resnet_model.pt";
    try
    {
        module = torch::jit::load(str);
    }
    catch (const c10::Error& e)
    {
        std::cerr << "12313";
        return -1;
    }

    // 創建一個輸入
    std::vector<torch::jit::IValue> inputs;
    inputs.push_back(torch::ones({ 1, 3, 224, 224 }));
    // 推理
    at::Tensor output = module.forward(inputs).toTensor();
    std::cout << output.slice(/*dim=*/1, /*start=*/0, /*end=*/5) << '\n';

    return 1;
}

最后匯總下：

python模型的序列化、保存代碼：

import torchvision
import torch

model = torchvision.models.resnet18()

example = torch.rand(1, 3, 224, 224)

traced_script_module = torch.jit.trace(model, example)

output = traced_script_module(torch.ones(1, 3, 224, 224))

#traced_script_module.save("traced_resnet_model.pt") # 和下面等價，格式名稱不同，僅此而已，在libtorch中是一樣的
traced_script_module.save("traced_resnet_model.torchscript.pt")
print()

libtorch的模型加載，推理代碼：

#include<torch/torch.h>
#include<torch/script.h>
#include<iostream>
#include<memory>

int main()
{
    torch::jit::script::Module module;
    std::string str = "traced_resnet_model.pt"; 
    //std::string str = "traced_resnet_model.torchscript.pt"; // 和上面等價，模型格式而已
    try
    {
        module = torch::jit::load(str);
    }
    catch (const c10::Error& e)
    {
        std::cerr << "12313";
        return -1;
    }

    // 創建一個輸入
    std::vector<torch::jit::IValue> inputs;
    inputs.push_back(torch::ones({ 1, 3, 224, 224 }));
    // 推理
    at::Tensor output = module.forward(inputs).toTensor();
    std::cout << output.slice(/*dim=*/1, /*start=*/0, /*end=*/5) << '\n';

    return 1;
}

 

 

reference：[1] https://pytorch.org/tutorials/advanced/cpp_export.html

PyTorch模型部署建議方案：

[2] https://blog.csdn.net/zhoumoon/article/details/104850615

[3] torch.jit.trace & torch.jit.script

https://www.dazhuanlan.com/2020/02/18/5e4b46eb099e3/