用 C++ 加载 TorchScript 模型（libtorch官方demo）

Q : Libtorch是什么及为什么用Libtorch？
A : Libtorch是pytorch的C++版本，现在的很多大型项目都是用C++写的，想使用训练好的模型，需要通过caffe等方式去集成，比较麻烦。这里pytorch官方提出了Libtorch，我们就可以把pytorch训练好的模型，打包起来，直接在C++工程中去用就好了，相比较caffe等，非常方便！

顾名思义，PyTorch 的主要接口是 Python 编程语言。虽然对于许多需要动态性和易于迭代的场景来说，Python 是一种合适且首选的语言，但同样在许多情况下，Python 的这些属性恰恰是不利的。后者经常应用的一种环境是生产——低延迟和严格部署要求的环境。对于生产场景，C++ 通常是首选语言，即使只是将其绑定到 Java、Rust 或 Go 等另一种语言。以下段落将概述 PyTorch 提供的从现有 Python 模型到可以纯粹从 C++ 加载和执行的序列化表示的路径，而不依赖于 Python。

---

第 1 步：将 PyTorch 模型转换为 Torch 脚本

PyTorch 模型从 Python 到 C++ 的过程是通过 Torch Script 实现的，Torch Script 是 PyTorch 模型的表示形式，可以由 Torch Script 编译器理解、编译和序列化。如果您从用普通“eager”API 编写的现有 PyTorch 模型开始，则必须首先将模型转换为 Torch 脚本。在下面讨论的最常见的情况下，这只需要很少的努力。

有两种将 PyTorch 模型转换为 Torch 脚本的方法。第一个称为tracing，这是一种通过使用示例输入对其进行一次评估并记录这些输入通过模型的流程来捕获模型结构的机制。这适用于限制使用控制流的模型。第二种方法是向模型添加显式注释，通知 Torch Script 编译器它可以直接解析和编译您的模型代码，但须遵守 Torch Script 语言施加的约束。

方法一：通过tracing转换为 Torch 脚本

要通过tracing将 PyTorch 模型转换为 Torch 脚本，您必须将模型的实例以及示例输入传递给 torch.jit.trace 函数。这将生成一个 torch.jit.ScriptModule 对象，其中模型评估的跟踪嵌入到模块的前向方法中：

import torch
import torchvision

# 自己模型的实例
model = torchvision.models.resnet18()

# 提供给模型的示例输入
example = torch.rand(1, 3, 224, 224)

# 使用 torch.jit.trace 通过tracing生成 torch.jit.ScriptModule。
traced_script_module = torch.jit.trace(model, example)

现在可以以与常规 PyTorch 模块相同的方式评估跟踪的 ScriptModule：

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>)

---

总结：

Tracing这种方法操作比较简单，只需要给模型一组输入，走一遍推理网络，然后由torch.jit.trace记录一下路径上的信息并保存即可。缺点是如果模型中存在控制流比如if-else语句，一组输入只能遍历一个分支，这种情况下就没办法完整的把模型信息记录下来。

---

方法二：通过Annotation注释转换为 Torch 脚本

在某些情况下，例如如果您的模型采用特定形式的控制流，您可能希望直接在 Torch Script 中编写模型并相应地注释您的模型。例如，假设您有以下普通 Pytorch 模型：

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

由于该模块的forward方法使用依赖于输入的控制流，因此不适合tracing。所以我们可以将其转换为 ScriptModule。为了将模块转换为 ScriptModule，需要使用 torch.jit.script 编译模块，如下所示：

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

---

总结：

直接在Torch脚本中编写模型并相应地注释模型，通过torch.jit.script编译模块，将其转换为ScriptModule。forward方法会被默认编译，forward中被调用的方法也会按照被调用的顺序被编译如果想要编译一个forward以外且未被forward调用的方法，可以添加 @torch.jit.export；如果你想在模型的forward方法中使用TorchScrip尚不支持的一些python特性，使用@torch.jit.ignore来修饰。

---

第 2 步：将脚本模块序列化为文件 

一旦您拥有了 ScriptModule（通过tracing或注释 PyTorch 模型），您就可以将其序列化为文件了。后面就能够使用 C++ 从该文件加载模块并执行它，而不依赖于 Python。假设我们想要序列化前面在跟踪示例中显示的 ResNet18 模型。要执行此序列化，只需调用模块上的 save 并向其传递一个文件名：

1.traced_script_module.save("traced_resnet_model.pt")
2.sm.save("my_module_model.pt") 

我宣布完成以上工作，现在已经正式离开了 Python 领域，并准备好跨入 C++ 领域。

第 3 步：在 C++ 中加载脚本模块

要在 C++ 中加载序列化的 PyTorch 模型，您的应用程序必须依赖于 PyTorch C++ API（也称为 LibTorch）。LibTorch 发行版包含共享库、头文件和 CMake 构建配置文件的集合。虽然 CMake 不是依赖 LibTorch 的必要条件，但它是推荐的方法，并且将来会得到很好的支持。在本教程中，我们将使用 CMake 和 LibTorch 构建一个最小的 C++ 应用程序，该应用程序只需加载并执行序列化的 PyTorch 模型。

#include <torch/script.h> // 导入torch/script.h头文件，其中包含了所需的API。

#include <iostream>
#include <memory>

int main(int argc, const char* argv[]) {
//检查命令行参数的数量是否正确。它要求命令行参数的数量必须是2，否则会输出错误消息并返回错误码-1。
  if (argc != 2) {
    std::cerr << "usage: example-app <path-to-exported-script-module>\n";
    return -1;
  }


  torch::jit::script::Module module;// 声明torch::jit::script::Module对象用于存储加载的脚本模块。
  try {
    // 使用torch::jit::load()从文件中反序列化脚本模块。
    module = torch::jit::load(argv[1]);
  }
  catch (const c10::Error& e) {
    std::cerr << "error loading the model\n";
    return -1;
  }

  std::cout << "ok\n"; // 打印"ok"表示成功加载和初始化模型。
}

假设我们将上述代码存储到名为 example-app.cpp 的文件中。构建它的最小 CMakeLists.txt 看起来很简单：

#设置了CMake的最小版本要求和项目名称。
cmake_minimum_required(VERSION 3.0 FATAL_ERROR)
project(custom_ops)
#查找Torch库
find_package(Torch REQUIRED)
#创建了一个名为 "example-app" 的可执行文件，并将源文件 "example-app.cpp" 添加到该可执行文件中。
add_executable(example-app example-app.cpp)
#将 Torch 库链接到 "example-app" 可执行文件。
target_link_libraries(example-app "${TORCH_LIBRARIES}")
#将以C++14标准进行编译源文件
set_property(TARGET example-app PROPERTY CXX_STANDARD 14)

在bash中执行编译命令

mkdir build
cd build
cmake -DCMAKE_PREFIX_PATH=/path/to/libtorch ..
cmake --build . --config Release

一切顺利的话你应该可以在build文件夹中看到你的程序的可执行文件了。

第 4 步：在 C++ 中执行脚本模块

在 C++ 中成功加载序列化的 ResNet18 后，我们现在只需几行代码即可执行它！让我们将这些行添加到 C++ 应用程序的 main() 函数中：

// 创建输入向量。
#std::vector是C++标准库中的容器，可以动态地调整大小以容纳多个元素。
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';
//这一行代码使用 slice 函数对 output 张量进行切片操作。slice 函数接受三个参数：dim 表示要切片的维度，start 表示起始索引，end 表示结束索引。在这个例子中，我们将维度1（通道维度）的部分切片出来，起始索引为0，结束索引为5。然后，使用 std::cout 将切片结果打印到标准输出流。

前两行设置模型的输入。我们创建一个 torch::jit::IValue 向量（类型擦除值类型 script::Module 方法接受和返回）并添加单个输入。为了创建输入张量，我们使用 torch::ones()，相当于 C++ API 中的 torch.ones。然后我们运行 script::Module 的forward 方法，将我们创建的输入向量传递给它。作为回报，我们得到一个新的 IValue，我们通过调用 toTensor() 将其转换为张量。

总结：本教程有望让您对 PyTorch 模型从 Python 到 C++ 的路径有一个大致的了解。通过本教程中描述的概念，您应该能够从普通的“eager”PyTorch 模型，到 Python 中编译的 ScriptModule，到磁盘上的序列化文件，以及（为了关闭循环）到可执行脚本： C++ 中的模块。当然，还有很多概念我们没有涉及到。例如，您可能会发现自己想要使用用 C++ 或 CUDA 实现的自定义运算符来扩展 ScriptModule，并在纯 C++ 生产环境中加载的 ScriptModule 内执行此自定义运算符。

参考资料：

https://pytorch.org/tutorials/advanced/cpp_export.html