libtorch的c++，加载*.pth

一、转换模型为TorchScript

前提：python只保存了参数，没存结构

要在C++中使用libtorch（PyTorch的C++接口），读取和加载通过torch.save保存的模型（ torch.save(pdn.state_dict()这种方式，只保存了参数，没存结构），需要转换模型为TorchScript。在python下实现。

复制代码

def get_pdn_small(out_channels=384, padding=False):
    pad_mult = 1 if padding else 0
    return nn.Sequential(
        nn.Conv2d(in_channels=3, out_channels=128, kernel_size=4,
                  padding=3 * pad_mult),
        nn.ReLU(inplace=True),
        nn.AvgPool2d(kernel_size=2, stride=2, padding=1 * pad_mult),
        nn.Conv2d(in_channels=128, out_channels=256, kernel_size=4,
                  padding=3 * pad_mult),
        nn.ReLU(inplace=True),
        nn.AvgPool2d(kernel_size=2, stride=2, padding=1 * pad_mult),
        nn.Conv2d(in_channels=256, out_channels=256, kernel_size=3,
                  padding=1 * pad_mult),
        nn.ReLU(inplace=True),
        nn.Conv2d(in_channels=256, out_channels=out_channels, kernel_size=4)
    )

def get_pdn_medium(out_channels=384, padding=False):
    pad_mult = 1 if padding else 0
    return nn.Sequential(
        nn.Conv2d(in_channels=3, out_channels=256, kernel_size=4,
                  padding=3 * pad_mult),
        nn.ReLU(inplace=True),
        nn.AvgPool2d(kernel_size=2, stride=2, padding=1 * pad_mult),
        nn.Conv2d(in_channels=256, out_channels=512, kernel_size=4,
                  padding=3 * pad_mult),
        nn.ReLU(inplace=True),
        nn.AvgPool2d(kernel_size=2, stride=2, padding=1 * pad_mult),
        nn.Conv2d(in_channels=512, out_channels=512, kernel_size=1),
        nn.ReLU(inplace=True),
        nn.Conv2d(in_channels=512, out_channels=512, kernel_size=3,
                  padding=1 * pad_mult),
        nn.ReLU(inplace=True),
        nn.Conv2d(in_channels=512, out_channels=out_channels, kernel_size=4),
        nn.ReLU(inplace=True),
        nn.Conv2d(in_channels=out_channels, out_channels=out_channels,
                  kernel_size=1)
    )

import torch

# 假设你有一个已训练的模型
model = get_pdn_small()

# 加载模型的state_dict
model.load_state_dict(torch.load('teacher_small.pth'))
model.eval()  # 设置模型为评估模式

# 将模型转化为TorchScript
scripted_model = torch.jit.script(model)
scripted_model.save('teacher_small.pt')

二、在C++中加载TorchScript模型：

在C++中，你可以使用torch::jit::load来加载.pt文件，如下所示：

复制代码

#include <torch/script.h>  // One-stop header for loading TorchScript models
#include <iostream>
#include <memory>

int main() {
    // 加载TorchScript模型
    try {
        // 加载模型
                std::shared_ptr<torch::jit::Module> model = std::make_shared<torch::jit::Module>(torch::jit::load("teacher_small.pt"));


        std::cout << "Model loaded successfully!" << std::endl;

        // 你可以在这里使用模型进行推理，比如输入一个张量
        // 例如，如果输入是一个3x224x224的图像，你需要创建一个相应的Tensor
        torch::Tensor input = torch::randn({1, 3, 224, 224});  // 示例输入
        std::vector<torch::jit::IValue> inputs;
        inputs.push_back(input);

        // 执行模型推理
        at::Tensor output = model->forward(inputs).toTensor();
        std::cout << "Output tensor: " << output << std::endl;
    }
    catch (const c10::Error& e) {
        std::cerr << "Error loading the model: " << e.what() << std::endl;
        return -1;
    }
}