Jetson 上的 TensorFlow/TensorRT 模型

此仓库包含用于在 NVIDIA Jetson 上使用 TensorFlow 图像分类和目标检测模型的脚本及文档。这些模型来源于 TensorFlow 模型仓库，并使用 TensorRT 进行优化。

• 设置
• 图像分类
  • 模型
  • 下载预训练模型
  • 构建 TensorRT / Jetson 兼容图
  • 使用 TensorRT 优化
  • Jupyter Notebook 示例
  • 针对自定义任务进行训练
• 目标检测
  • 模型
  • 下载预训练模型
  • 构建 TensorRT / Jetson 兼容图
  • 使用 TensorRT 优化
  • Jupyter Notebook 示例
  • 针对自定义任务进行训练

设置

1. 使用 JetPack 3.2（包含 TensorRT）为您的 Jetson TX2 刷机。

2. 在 Jetson 上安装其他依赖项：

   sudo apt-get install python-pip python-matplotlib python-pil
   

3. 安装 TensorFlow 1.7+（支持 TensorRT）。下载 预编译的 pip 轮子 并使用 pip 安装。

   pip install tensorflow-1.8.0-cp27-cp27mu-linux_aarch64.whl --user
   

   或者如果您使用的是 Python 3：

   pip3 install tensorflow-1.8.0-cp35-cp35m-linux_aarch64.whl --user
   

4. 克隆此仓库：

   git clone --recursive https://github.com/NVIDIA-Jetson/tf_trt_models.git
   cd tf_trt_models
   

5. 运行安装脚本：

   ./install.sh
   

   或者如果您想指定 Python 解释器：

   ./install.sh python3
   

图像分类

模型

模型	输入尺寸	TF-TRT TX2	TF TX2
inception_v1	224x224	7.36ms	22.9ms
inception_v2	224x224	9.08ms	31.8ms
inception_v3	299x299	20.7ms	74.3ms
inception_v4	299x299	38.5ms	129ms
inception_resnet_v2	299x299		158ms
resnet_v1_50	224x224	12.5ms	55.1ms
resnet_v1_101	224x224	20.6ms	91.0ms
resnet_v1_152	224x224	28.9ms	124ms
resnet_v2_50	299x299	26.5ms	73.4ms
resnet_v2_101	299x299	46.9ms
resnet_v2_152	299x299	69.0ms
mobilenet_v1_0p25_128	128x128	3.72ms	7.99ms
mobilenet_v1_0p5_160	160x160	4.47ms	8.69ms
mobilenet_v1_1p0_224	224x224	11.1ms	17.3ms

TF - 原始 TensorFlow 图（FP32）

TF-TRT - TensorRT 优化后的图（FP16）

以上基准测试时间是在将 Jetson TX2 设置为 MAX-N 模式后收集的。为此，请在终端中运行以下命令：

sudo nvpmodel -m 0
sudo ~/jetson_clocks.sh

下载预训练模型

为了方便起见，我们提供了一个脚本，用于从 TensorFlow 模型仓库下载预训练模型。

from tf_trt_models.classification import download_classification_checkpoint

checkpoint_path = download_classification_checkpoint('inception_v2')

要手动下载预训练模型，请按照 此处 的链接操作。

构建 TensorRT / Jetson 兼容图

from tf_trt_models.classification import build_classification_graph

frozen_graph, input_names, output_names = build_classification_graph(
    model='inception_v2',
    checkpoint=checkpoint_path,
    num_classes=1001
)

使用 TensorRT 优化

import tensorflow.contrib.tensorrt as trt

trt_graph = trt.create_inference_graph(
    input_graph_def=frozen_graph,
    outputs=output_names,
    max_batch_size=1,
    max_workspace_size_bytes=1 << 25,
    precision_mode='FP16',
    minimum_segment_size=50
)

Jupyter Notebook 示例

有关执行上述步骤并在实际图像上运行的完整示例，请参阅 jupyter notebook 示例。

针对自定义任务进行训练

请遵循 TensorFlow 模型仓库 中的文档。获得检查点后，按照上述方法构建图并使用 TensorRT 进行优化。

目标检测

模型

模型	输入尺寸	TF-TRT TX2	TF TX2
ssd_mobilenet_v1_coco	300x300	50.5ms	72.9ms
ssd_inception_v2_coco	300x300	54.4ms	132ms

TF - 原始 TensorFlow 图（FP32）

TF-TRT - TensorRT 优化后的图（FP16）

以上基准测试时间是在将 Jetson TX2 设置为 MAX-N 模式后收集的。为此，请在终端中运行以下命令：

sudo nvpmodel -m 0
sudo ~/jetson_clocks.sh

下载预训练模型

为了方便起见，我们提供了一个脚本，用于从 TensorFlow 模型仓库下载预训练模型权重和配置文件。

from tf_trt_models.detection import download_detection_model

config_path, checkpoint_path = download_detection_model('ssd_inception_v2_coco')

要手动下载预训练模型，请按照 此处 的链接操作。

重要提示： 一些目标检测配置文件具有非常低的非极大值抑制阈值（例如 1e-8）。这可能会导致不必要的大量 CPU 后处理负载。根据您的应用，建议将此值提高到更大的数值（如 0.3），以提升性能。我们在上述基准测试中就是这么做的。您可以在调用 build_detection_graph 之前直接修改配置文件来实现这一点。该参数例如可在这一 行 找到。

构建 TensorRT / Jetson 兼容图

from tf_trt_models.detection import build_detection_graph

frozen_graph, input_names, output_names = build_detection_graph(
    config=config_path,
    checkpoint=checkpoint_path
)

使用 TensorRT 进行优化

import tensorflow.contrib.tensorrt as trt

trt_graph = trt.create_inference_graph(
    input_graph_def=frozen_graph,
    outputs=output_names,
    max_batch_size=1,
    max_workspace_size_bytes=1 << 25,
    precision_mode='FP16',
    minimum_segment_size=50
)

Jupyter Notebook 示例

有关执行上述步骤并在实际图像上运行的完整示例，请参阅 Jupyter Notebook 示例。

针对自定义任务进行训练

请按照 TensorFlow 模型仓库 中的文档操作。在获得检查点后，按照上述方法构建图并使用 TensorRT 进行优化。请注意，并非所有模型都经过测试，因此在训练时应使用与 ssd_mobilenet_v1_coco 或 ssd_inception_v2_coco 模型类似的物体检测配置文件。部分配置参数（如类别数、图像尺寸、非极大值抑制参数等）可以修改，但性能可能会有所不同。

tf_trt_models 快速上手指南

本指南帮助开发者在 NVIDIA Jetson 平台上快速部署经过 TensorRT 优化的 TensorFlow 图像分类与目标检测模型。

环境准备

• 硬件平台：NVIDIA Jetson TX2 (或其他支持 JetPack 的 Jetson 设备)
• 系统镜像：需刷入 JetPack 3.2 (必须包含 TensorRT 组件)
• 性能模式：建议将设备设置为最大性能模式以获得最佳推理速度：

  sudo nvpmodel -m 0
  sudo ~/jetson_clocks.sh
  

安装步骤

1. 安装系统依赖

   sudo apt-get install python-pip python-matplotlib python-pil
   

2. 安装 TensorFlow (含 TensorRT 支持) 下载适用于 Jetson 的预编译 wheel 包（TensorFlow 1.8+），然后使用 pip 安装。

   Python 2.7:

   pip install tensorflow-1.8.0-cp27-cp27mu-linux_aarch64.whl --user
   

   Python 3.5:

   pip3 install tensorflow-1.8.0-cp35-cp35m-linux_aarch64.whl --user
   

   提示：国内用户若下载官方 wheel 包较慢，可尝试在 NVIDIA 开发者论坛查找国内镜像资源或使用本地传输方式。

3. 克隆项目并安装

   git clone --recursive https://github.com/NVIDIA-Jetson/tf_trt_models.git
   cd tf_trt_models
   ./install.sh
   

   若需指定 Python 3 环境：

   ./install.sh python3
   

基本使用

1. 图像分类 (Image Classification)

以下示例展示如何下载预训练模型、构建图并进行 TensorRT 优化（以 inception_v2 为例）。

from tf_trt_models.classification import download_classification_checkpoint, build_classification_graph
import tensorflow.contrib.tensorrt as trt

# 1. 下载预训练模型
checkpoint_path = download_classification_checkpoint('inception_v2')

# 2. 构建兼容 TensorRT/Jetson 的冻结图
frozen_graph, input_names, output_names = build_classification_graph(
    model='inception_v2',
    checkpoint=checkpoint_path,
    num_classes=1001
)

# 3. 使用 TensorRT 优化 (FP16 精度)
trt_graph = trt.create_inference_graph(
    input_graph_def=frozen_graph,
    outputs=output_names,
    max_batch_size=1,
    max_workspace_size_bytes=1 << 25,
    precision_mode='FP16',
    minimum_segment_size=50
)

完整示例：参考 examples/classification/classification.ipynb。

2. 目标检测 (Object Detection)

以下示例展示如何配置并优化目标检测模型（以 ssd_inception_v2_coco 为例）。

from tf_trt_models.detection import download_detection_model, build_detection_graph
import tensorflow.contrib.tensorrt as trt

# 1. 下载模型权重与配置文件
config_path, checkpoint_path = download_detection_model('ssd_inception_v2_coco')

# 2. 构建冻结图
frozen_graph, input_names, output_names = build_detection_graph(
    config=config_path,
    checkpoint=checkpoint_path
)

# 3. 使用 TensorRT 优化
trt_graph = trt.create_inference_graph(
    input_graph_def=frozen_graph,
    outputs=output_names,
    max_batch_size=1,
    max_workspace_size_bytes=1 << 25,
    precision_mode='FP16',
    minimum_segment_size=50
)

性能优化提示：部分默认配置文件的非极大值抑制（NMS）阈值过低（如 1e-8），会导致 CPU 后处理负载过高。建议在调用 build_detection_graph 前修改配置文件，将阈值调整为合理值（如 0.3）以提升整体性能。 完整示例：参考 examples/detection/detection.ipynb。