在C++中加载TorchScript模型的方法

在C++中加载TorchScript模型的方法

如果我们想要在C++中加载TorchScript模型（.pt或.pkl文件），则需要使用到libtorch库和TorchScript API。下面是加载模型的完整攻略：

1. 下载libtorch库

在pytorch官网下载适合自己操作系统的libtorch库，解压后即可得到所需的头文件和库文件。

1. 编写加载模型的代码

2.1 加载模型

首先需要按以下方式调用torch::jit::load函数加载模型：

torch::jit::script::Module module;
try {
    module = torch::jit::load(model_path);
} catch (const c10::Error& e) {
    std::cerr << "Error loading the model\n";
    return -1;
}

其中，model_path是模型的路径。

2.2 获取输入和输出的名称

在编写前向推理代码之前，需要先知道我们模型的输入和输出张量名。可以通过以下代码获取：

std::vector<std::string> input_names;
std::vector<std::string> output_names;
for (const auto& param : module.named_parameters()) {
    input_names.push_back(param.name);
}
for (const auto& sub_module : module.named_modules()) {
    for (const auto& param : sub_module.value.named_parameters()) {
        input_names.push_back(param.name);
    }
    for (const auto& buffer : sub_module.value.named_buffers()) {
        input_names.push_back(buffer.name);
    }
}
for (const auto& output : module.get_output_nodes()) {
    for (const auto& output_val : output->outputs()) {
        output_names.push_back(output_val->debugName());
    }
}

2.3 构建输入张量

构建需要传入模型的输入张量，需要先定义一个std::vector类型的变量inputs，将需要传入模型的所有输入张量通过该vector传递：

std::vector<torch::jit::IValue> inputs;

// 第一个张量作为输入
inputs.push_back(torch::ones({1, 3, 224, 224}));

2.4 前向推理

前向推理的代码如下：

// 将模型移动到eval模式
module.eval();

// 前向推理，返回一个torch::jit::IValue类型的结果
at::Tensor output = module.forward(inputs).toTensor();

前向推理的输出将由output变量保存，你可以通过它来获取模型的输出张量。

2.5 输出结果

前向推理完成后，我们需要将输出结果转换为可读的形式。无论是标量、张量、矩阵或其他数据，都可以通过以下代码获取：

// 输出结果
std::cout << output.slice(/*dim=*/1, /*start=*/0, /*end=*/5) << '\n';

至此，我们已经成功地从C++代码中加载TorchScript模型并进行了前向推理。

1. 示例说明

以下是两个示例，分别演示如何加载经过训练的分类网络和循环神经网络模型。

示例1：加载分类模型

在这个示例中，我们假设我们已经训练好了一个分类网络，并将其保存为.pt或.pkl文件。现在我们想要在C++中加载该模型，以便进行分类任务的前向推理。

首先，我们需要修改加载模型的代码，以适应我们的网络：

// 加载模型
torch::jit::script::Module module;
try {
    module = torch::jit::load(model_path);
} catch (const c10::Error& e) {
    std::cerr << "Error loading the model\n";
    return -1;
}

// 获取输入和输出的名称
std::vector<std::string> input_names;
std::vector<std::string> output_names;
input_names.push_back("input");
output_names.push_back("output");

// 构建输入张量
std::vector<torch::jit::IValue> inputs;
inputs.push_back(torch::rand({1, 3, 224, 224}));

// 前向推理
module.eval();
at::Tensor output = module.forward(inputs).toTensor();

// 输出结果
std::cout << output.slice(/*dim=*/1, /*start=*/0, /*end=*/5) << '\n';

示例2：加载循环神经网络模型

在这个示例中，我们加载的是一个循环神经网络模型，用于执行语言建模任务。假设我们已经通过训练获得了一个.pth文件，并将其保存在了model.pt路径下。我们需要加载该模型，并在输入一个序列后计算模型输出。

加载模型的代码与前面的示例相同：

torch::jit::script::Module module;
try {
    module = torch::jit::load("model.pt");
} catch (const c10::Error& e) {
    std::cerr << "error loading the model\n";
    return -1;
}

获取输入和输出的名称与前面的示例也类似：

std::vector<std::string> input_names;
std::vector<std::string> output_names;
input_names.push_back("input");
output_names.push_back("output");

在这个示例中，我们需要输入一个序列，但是C++不支持Python中的特殊对象（如List、Tuple等），因此我们需要使用张量表示序列。张量的shape必须是[N, L]，其中N是序列中的单词数，而L是最大句子长度。我们可以使用以下代码将字符串序列转换为张量：

std::vector<std::string> words = {"hello", "world"};
std::vector<int64_t> data;
for (const auto& word : words) {
    // 使用一个简单的哈希函数将字符串映射为整数
    data.push_back(static_cast<int64_t>(std::hash<std::string>()(word)));
}
auto input = torch::tensor(data).reshape({1, data.size()});

构建输入张量后，我们可以执行前向推理：

std::vector<torch::jit::IValue> inputs;
inputs.push_back(input);
module.eval();
auto output = module.forward(inputs).toTensor();

前向推理完成后，我们需要将输出结果解码为可读的形式：

std::cout << output.slice(/*dim=*/1, /*start=*/0, /*end=*/5) << '\n';

以上就是两个示例，它们分别展示了如何加载TorchScript中的分类网络和循环神经网络模型。你可以根据自己的需求自由地修改这些示例，以便于适应不同的模型和输入。