要在C++中调用PyTorch模型,可以使用LibTorch库。以下是一个简单的示例代码,演示了如何加载一个PyTorch模型并使用输入数据进行推理:
#include <torch/torch.h>
#include <iostream>
int main() {
// 加载模型
torch::jit::script::Module module;
try {
module = torch::jit::load("path/to/model.pt");
} catch (const c10::Error& e) {
std::cerr << "Error loading the model\n";
return -1;
}
// 准备输入数据
torch::Tensor input = torch::ones({1, 3, 224, 224}); // 示例输入数据
// 运行推理
at::Tensor output = module.forward({input}).toTensor();
// 输出结果
std::cout << "Output tensor: " << output << std::endl;
return 0;
}
在这个示例中,首先加载了一个PyTorch模型(假设模型保存在model.pt
文件中)。然后创建了一个示例输入张量input
,并将其传递给模型进行推理。最后,输出了模型的输出张量。
请注意,为了能够编译这段代码,需要在项目中链接LibTorch库并设置正确的包含路径。更多关于LibTorch的用法和配置信息,请参考PyTorch官方文档。