c++怎么调用LibTorch运行深度学习模型_c++ PyTorch模型加载与推理【案例】

C++加载LibTorch模型需用torch::jit::load()加载.pt脚本模型(非.pth),调用module.eval()和module.to()统一设备,输入张量须严格匹配shape/dtype/device,forward()返回IValue需按类型(toTensor/toGenericDict/toTuple)安全解包,GPU结果须.to(kCPU)后访问。

怎么用 C++ 加载 torch::jit::script::Module 并运行推理

LibTorch 的核心是 torch::jit::script::Module,它对应 Python 侧用 torch.jit.tracetorch.jit.script 导出的 .pt 模型文件。C++ 里不能直接加载 .pth(state dict),必须用脚本模型格式。

常见错误:把训练时保存的 model.pth(含 state_dict)直接传给 torch::jit::load(),会报错 Expected a script module, but got a state dict

  • Python 侧导出必须用 traced_model.save("model.pt")scripted_model.save("model.pt")
  • C++ 侧加载用 torch::jit::load("model.pt"),返回 torch::jit::script::Module
  • 加载后建议调用 module.to(torch::kCUDA)(如果用 GPU)和 module.eval(),否则可能因 dropout/batch norm 行为异常导致输出不一致

输入张量怎么构造才匹配模型期望的 forward 签名

模型在 Python 中 trace/script 时记录了输入 shape、dtype 和名称(如果是命名 tuple 或 dict)。C++ 侧必须严格对齐,否则运行时报 Expected object of scalar type Float but got scalar type LongWrong number of arguments

典型场景:图像分类模型输入是 [1, 3, 224, 224]float32 张量,但你传了 int8[3, 224, 224](缺 batch 维)。

  • torch::ones({1, 3, 22

    4, 224}, torch::kFloat)
    构造占位输入,注意 dtype 和 device 要与模型一致
  • 图像预处理(归一化、resize)必须在 C++ 里重做,不能依赖 Python 的 transforms;常用 OpenCV 读图后转 torch::from_blob + permute + to
  • 如果模型接受多个输入(如 forward(x, mask)),需用 std::vector<:jit::ivalue> 包装,顺序不能错

为什么 module.forward() 返回值取不到数据?

module.forward() 返回的是 torch::jit::IValue,不是 torch::Tensor。直接打印或取 .data_ptr() 会崩溃或得到乱码。

常见错误现象:程序 segfault,或输出 IValue of None,或 Expected Tensor but got GenericDict(模型返回字典时没解包)。

  • 先用 .toTensor() 尝试转成 Tensor(适用于单输出)
  • 若模型返回字典(如 DETR),用 .toGenericDict(),再用 .at("pred_logits") 取 key
  • 若返回 tuple,用 .toTuple()->elements() 拆包
  • 取值前务必检查 ivalue.isTensor() 等类型断言,避免运行时异常

GPU 推理卡死或结果全零?检查这三处

LibTorch 的 CUDA 支持不像 Python 那样“自动兜底”,很多细节要手动对齐。

  • 编译 LibTorch 时必须选 CUDA 版本(下载页明确标有 cu118cu121),且与系统 nvidia-drivercudnn 兼容;混用 cpu-only 库和 .to(kCUDA) 会静默失败
  • 所有输入 Tensor、模型、中间变量必须在同一个 device 上:创建时加 .to(torch::kCUDA),不要只移模型不移输入
  • GPU 推理后,结果 Tensor 仍位于 GPU,需调用 .to(torch::kCPU).detach().clone() 再访问数据,否则 .data_ptr() 指向显存地址,CPU 不可读
#include 
#include 

int main() {
    auto module = torch::jit::load("resnet18.pt");
    module.to(torch::kCUDA);
    module.eval();

    // 构造输入:[1,3,224,224] float32 GPU tensor
    auto input = torch::randn({1, 3, 224, 224}, torch::kFloat)
                     .to(torch::kCUDA);

    // 推理
    at::IValue output = module.forward({input});
    auto out_tensor = output.toTensor().to(torch::kCPU).detach().clone();

    // 打印 top-1 index
    std::cout << out_tensor.argmax(1) << std::endl;
}
LibTorch 的坑大多藏在类型转换和设备一致性上,而不是 API 调用本身。一个 to(kCUDA) 漏写,或一个 .toTensor() 忘加,就足以让程序跑出错、结果错、或者干脆不报错但输出全零。