一、转换模型为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 0return 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 0return 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的图像，你需要创建一个相应的Tensortorch::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;}
}
)
