PyTorch 在 C ++中加载 TorchScript 模型

PyTorch 在 C++ 中加载 TorchScript 模型教程

将 PyTorch 模型从 Python 环境成功移植到 C++ 环境，能够充分发挥 C++ 在高性能计算和生产部署中的优势，同时保留 PyTorch 模型的灵活性和强大功能。本文将详细指导您完成这一过程，包括模型转换、序列化、C++ 应用程序开发以及模型加载与执行等关键步骤。

一、从 PyTorch 模型到 TorchScript 的转换

（一）通过追踪进行转换

对于大多数模型，尤其是控制流较为简单的模型，使用追踪方法可以轻松地将其转换为 TorchScript 格式：

import torch
import torchvision


## 定义模型
model = torchvision.models.resnet18()


## 准备示例输入
example_input = torch.rand(1, 3, 224, 224)


## 使用追踪生成 TorchScript 模型
traced_script_module = torch.jit.trace(model, example_input)

（二）通过注释进行转换

若模型包含复杂的控制流，需要使用注释方法进行转换：

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


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

二、TorchScript 模型的序列化

将生成的 TorchScript 模型保存到文件，以便在 C++ 应用程序中加载和使用：

## 序列化追踪得到的模型
traced_script_module.save("traced_resnet_model.pt")


## 序列化注释得到的模型
scripted_module.save("scripted_my_module.pt")

三、在 C++ 中加载和执行 TorchScript 模型

（一）创建 C++ 应用程序

编写一个简单的 C++ 程序来加载和执行 TorchScript 模型：

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


int main(int argc, const char* argv[]) {
    if (argc != 2) {
        std::cerr << "Usage: " << argv[0] << " <path-to-model>\n";
        return -1;
    }


    torch::jit::script::Module module;
    try {
        module = torch::jit::load(argv[1]);
    } catch (const c10::Error& e) {
        std::cerr << "Error loading the model: " << e.what() << std::endl;
        return -1;
    }


    std::cout << "Model loaded successfully!\n";


    // 准备输入数据
    std::vector<torch::jit::IValue> inputs;
    inputs.push_back(torch::ones({1, 3, 224, 224}));


    // 执行模型
    at::Tensor output = module.forward(inputs).toTensor();


    // 输出结果
    std::cout << "Model output:\n" << output.slice(1, 0, 5) << std::endl;


    return 0;
}

（二）构建 C++ 应用程序

使用 CMake 构建上述 C++ 应用程序：

1. 创建 CMakeLists.txt 文件

cmake_minimum_required(VERSION 3.0 FATAL_ERROR)
project(PyTorchCppExample)


find_package(Torch REQUIRED)


add_executable(example-app example-app.cpp)


target_link_libraries(example-app "${TORCH_LIBRARIES}")


set_property(TARGET example-app PROPERTY CXX_STANDARD 14)

2. 构建命令

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

（三）运行 C++ 应用程序

将序列化后的模型文件路径作为参数传递给构建好的可执行文件：

./example-app traced_resnet_model.pt

预期输出示例：

Model loaded successfully!
Model output:
-0.2698 -0.0381  0.4023 -0.3010 -0.0448
[ Variable[CPUFloatType]{1,5} ]

四、提升与优化

（一）模型优化

在加载模型后，可以对模型进行进一步优化以提升性能：

module.to(torch::kCUDA); // 将模型移动到 GPU
inputs[0] = inputs[0].to(torch::kCUDA); // 确保输入数据也在 GPU 上

（二）自定义运算符

如果需要在 C++ 中实现自定义运算符，可以参考 PyTorch C++ API 文档进行开发和集成。

（三）多线程与并行处理

利用 C++ 的多线程库和 PyTorch 提供的并行计算功能，进一步提升模型推理速度。

（四）模型量化

对模型进行量化处理，减少模型大小并加速推理过程，适合在资源受限的环境中部署。

五、总结与拓展

通过本文，您已经学习了如何将 PyTorch 模型转换为 TorchScript 格式，并在 C++ 应用程序中加载和执行。这一过程使您能够在高性能、低延迟的生产环境中充分利用 PyTorch 模型的强大功能。未来，您可以探索更多高级特性，如在 C++ 中实现自定义运算符、优化模型性能以及与其他 C++ 框架和库进行集成等。编程狮将持续为您更新更多深度学习跨语言部署的实用教程，助力您的技术成长。

PyTorch 在 C ++中加载 TorchScript 模型