如何将pytorch模型转换为TensorRT？

Question

我已经在 Nvidia GPU 上训练了分类模型并保存了模型权重（checkpoint.pth）。如果我想在jetson nano中部署这个模型并测试它。

我应该将其转换为 TenorRT 吗？如何将其转换为TensorRT？

我对此很陌生。如果有人可以纠正我，那将会很有帮助。

Answer 1

实现该方式的最佳方式是从Pytorch导出Onnx模型。 接下来，使用TensorRT官方包提供的TensorRT工具

trtexec

，将onnx模型转换为TensorRT模型。

您可以参考此页面：https://github.com/NVIDIA/TensorRT/blob/master/samples/opensource/trtexec/README.md

TRTEXEC

是一个更原生的工具，您可以从NVIDIA NGC图像中获取它或直接从官方网站下载。

如果你使用torch2trt这样的工具，很容易遇到算子问题，而且解决起来确实很复杂（如果你不熟悉处理插件问题）。

Answer 2

您可以使用这个工具：

https://github.com/NVIDIA-AI-IOT/torch2trt

以下是如何将转换器实现到引擎文件的更多详细信息：

https://github.com/NVIDIA-AI-IOT/torch2trt/issues/254

Answer 3

您可以使用Torch-TensorRT。

Torch-TensorRT 是 PyTorch 的推理编译器，针对 NVIDIA 通过 NVIDIA 的 TensorRT 深度学习优化器和运行时的 GPU。它 通过以下方式支持即时 (JIT) 编译工作流程 torch.compile 接口以及提前 (AOT) 工作流程。 Torch-TensorRT 无缝集成到 PyTorch 生态系统中 支持优化 TensorRT 代码与标准的混合执行 PyTorch 代码。