dustynv/torch2trt

torch2trt 镜像文档

镜像概述和主要用途

torch2trt 是一个专为NVIDIA Jetson平台设计的Docker镜像，提供将PyTorch模型转换为TensorRT格式的工具链。通过该镜像，用户可在Jetson设备上便捷地进行模型优化，将PyTorch模型转换为TensorRT引擎，从而显著提升模型推理性能，适用于边缘计算场景下的深度学习应用部署。

核心功能和特性

支持版本

• L4T（Linux for Tegra）版本：>=32.6（对应JetPack 4.6及以上版本）
• 兼容JetPack版本：
  • L4T R32.7.x 容器：支持JetPack 4.6+
  • L4T R35.x 容器：支持JetPack 5.1+

核心依赖项

• 基础工具：build-essential、python、cmake
• 计算框架：cuda、cudnn、pytorch:2.2、torchvision
• 模型转换：onnx、tensorrt

镜像信息

使用场景和适用范围

• 边缘AI部署：在Jetson Nano、Jetson Xavier NX、Jetson AGX Orin等设备上部署PyTorch模型，通过TensorRT优化提升推理速度。
• 模型转换工作流：作为PyTorch到TensorRT转换的标准化环境，确保依赖兼容性，简化部署流程。
• 深度学习应用开发：支持计算机视觉、自然语言处理等领域的模型优化，适用于nanodb、nanoowl、nanosam、xtts等依赖项目。

使用方法和配置说明

运行容器

方法1：使用jetson-containers工具（推荐）

jetson-containers工具可自动匹配兼容的镜像版本，简化运行流程：

bash
# 自动拉取或构建兼容镜像并运行
jetson-containers run $(autotag torch2trt)

# 显式指定镜像标签
jetson-containers run dustynv/torch2trt:r35.2.1

方法2：直接使用docker run命令

手动指定镜像及运行参数：

bash
sudo docker run --runtime nvidia -it --rm --network=host dustynv/torch2trt:r35.2.1

挂载本地目录

通过-v参数将主机目录挂载到容器中，方便模型文件和数据共享：

bash
jetson-containers run -v /host/path:/container/path $(autotag torch2trt)

运行指定命令

在容器中直接执行命令（非交互式shell）：

bash
jetson-containers run $(autotag torch2trt) python3 convert_model.py --input model.pth --output model.trt

构建容器

若需手动构建镜像，先完成https://github.com/dusty-nv/jetson-containers/tree/master/docs/setup.md%EF%BC%8C%E5%86%8D%E6%89%A7%E8%A1%8C%EF%BC%9A

bash
jetson-containers build torch2trt

构建过程会自动集成所有依赖项并进行测试。使用--help查看更多构建选项：

bash
jetson-containers build torch2trt --help

依赖项目

该镜像为以下项目提供基础支持：

• l4t-pytorch
• nanodb
• nanoowl
• nanosam
• xtts