libtorch加载TorchScript模型必须用torch::jit::load(),仅支持torch.jit.script()或torch.jit.trace()导出的.pt/.pth文件;加载state_dict.pth会报“Expected a script module”错误。

libtorch 加载 torchscript 模型必须用 torch::jit::load()
不是 torch::load(),也不是 Python 里那种 torch.load(..., map_location=...) 的写法。C++ libtorch 只认序列化后的 TorchScript 模型(即 .pt 或 .pth 文件,但必须是通过 torch.jit.script() 或 torch.jit.trace() 导出的)。直接拿 PyTorch 的 state_dict.pth 会崩溃,报错类似:Expected a script module, but got a state dict。
实操建议:
立即学习“C++免费学习笔记(深入)”;
- Python 端导出模型时,务必用
traced_model.save("model.pt")或scripted_model.save("model.pt"),别用torch.save(model.state_dict(), ...) - C++ 端加载前加异常捕获:
try { module = torch::jit::load("model.pt"); } catch (const c10::Error& e) { std::cerr << "Failed to load model: " << e.msg() << std::endl; return -1; } - 路径必须是绝对路径或确保工作目录正确;相对路径容易因执行位置不同而失败
输入 tensor 必须匹配模型期望的 shape 和 dtype
PyTorch 模型在 trace/script 时“记住”的是某次前向的输入特征,比如 torch.randn(1, 3, 224, 224)。C++ 里如果喂进去 torch::randn({1, 3, 256, 256}),可能不报错但输出错乱;更常见的是 dtype 不匹配:Python 默认 float32,但 C++ 里 torch::randn({1,3,224,224}) 默认也是 float32,可如果你手动用了 torch::kFloat64,就会触发 Expected object of scalar type Float but got scalar type Double。
实操建议:
立即学习“C++免费学习笔记(深入)”;
- 用
module->get_method("forward").graph()或先在 Python 里打印traced_model.graph查看输入签名 - 构造输入时显式指定 dtype 和 device:
auto input = torch::randn({1, 3, 224, 224}, torch::kFloat).to(torch::kCUDA); // 如果模型在 GPU 上 - 注意 batch 维必须存在——即使只推一个样本,也要是
{1, ...},不能是{3, 224, 224}
GPU 推理需手动管理 device 和 memory
libtorch 不会自动把模型和输入放到同一设备。常见错误是模型在 CPU、输入在 CUDA,或者反过来,报错:Expected all tensors to be on the same device。而且 torch::jit::load() 默认加载到 CPU,即使你后续调用 .to(torch::kCUDA),也得确保 CUDA 初始化已完成(比如调用过 torch::cuda::is_available())。
实操建议:
立即学习“C++免费学习笔记(深入)”;
- 加载后立刻移到目标设备:
module = torch::jit::load("model.pt"); module->to(torch::kCUDA); - 输入 tensor 也要同步 move:
input = input.to(torch::kCUDA); - 别忘了 warmup:首次 CUDA 推理有显存分配开销,建议先跑一次 dummy 输入再计时
- 显存不会自动释放——
module和中间 tensor 都要靠 RAII 或手动.reset(),避免长期持有导致 OOM
多线程推理要小心 torch::jit::script::Module 的共享问题
同一个 torch::jit::script::Module 实例不能被多个线程同时调用 forward()。这不是线程安全的,会出现 segfault 或数值错误。libtorch 官方明确说:“Modules are not thread-safe for mutation”,而 forward() 内部可能修改内部状态(比如 BatchNorm 的 running stats,哪怕 eval 模式下某些版本也有风险)。
实操建议:
立即学习“C++免费学习笔记(深入)”;
- 每个线程持有一个独立的
Module实例(即重复调用torch::jit::load()),虽然耗内存,但最稳妥 - 如果必须共享,加全局 mutex 包裹
module.forward()调用,但会串行化,失去并发意义 - 注意
torch::set_num_threads(1)对 OpenMP 线程有用,但对模型 forward 并发无效——那是模块实例层的问题


















