Windows下Libtorch使用Torch TensorRT加载模型报错求助

问题根源 #

你通过Dynamo IR导出的TorchScript模块依赖Torch-TensorRT定义的自定义__torch__.torch.classes.tensorrt.Engine类，但C++的LibTorch环境默认不会自动注册该扩展类，导致加载模型时无法识别类型。

解决方案 #

1. 链接Torch-TensorRT的C++库 #

确保你的C++项目正确链接编译生成的Torch-TensorRT库文件：

• Windows下需要链接torchtrt.lib，并将torchtrt.dll放在可执行文件同目录或系统PATH中。
• 若使用CMake，添加以下配置：

find_package(Torch REQUIRED)
find_package(TorchTensorRT REQUIRED)
target_link_libraries(你的项目名称 PRIVATE ${TORCH_LIBRARIES} ${TORCHTENSORRT_LIBRARIES})

2. 显式初始化Torch-TensorRT扩展 #

在加载模型前，必须调用Torch-TensorRT的初始化函数注册自定义类，修改C++代码如下：

#include <torch/torch.h>
#include <torch_tensorrt/torch_tensorrt.h>

int main() {
    // 初始化Torch-TensorRT，注册自定义TensorRT Engine类
    torch_tensorrt::init();

    // 加载TorchScript模型
    auto trt_mod = torch::jit::load("trt.ts");

    // 后续推理逻辑...
    return 0;
}

3. 验证版本一致性 #

• 确认Python端Torch-TensorRT（2.4.0）与C++编译的Torch-TensorRT版本完全一致。
• 确保LibTorch（2.5.0.dev）和PyTorch（2.4.0+cu124）的API兼容（nightly版本需注意无破坏性变更）。

4. 备选：切换到TorchScript IR导出 #

若Dynamo方式仍存在兼容性问题，可改用TorchScript IR编译导出，兼容性更稳定：

import torch
import torch_tensorrt

model = MyModel().eval().cuda()
inputs = [torch.randn((1, 3, 224, 224)).cuda()]

# 使用torchscript IR编译
trt_gm = torch_tensorrt.compile(model, ir="torchscript", inputs=inputs, enabled_precisions={torch.float32})
torch_tensorrt.save(trt_gm, "trt_ts_script.ts", output_format="torchscript", inputs=inputs)

C++端仍需执行步骤1和2加载模型。

内容的提问来源于stack exchange，提问作者Mmmmyy