当前位置：首页 > news >正文

使用 TensorFlow SSD 网络进行对象检测

news 2026/2/10 4:37:04

使用 TensorFlow SSD 网络进行对象检测

描述
这个示例是如何工作的？
- 处理输入图
- 准备数据
- sampleUffSSD 插件
- 验证输出
- TensorRT API 层和操作
先决条件
运行示例
- 示例 --help 选项
附加资源
许可证
更改日志
已知问题

描述

该示例 sampleUffSSD 预处理 TensorFlow SSD 网络，使用 TensorRT 进行 SSD 网络的推理，利用 TensorRT 插件加速推理。

该示例基于 SSD：单次多盒检测器论文。SSD 网络在网络的单次前向传播中执行对象检测和定位任务。

本示例中使用的 SSD 网络基于 TensorFlow 实现的 SSD 网络，与原始论文有所不同，它具有 inception_v2 骨干网络。有关实际模型的更多信息，请下载 ssd_inception_v2_coco。TensorFlow SSD 网络是使用 MSCOCO 数据集对 InceptionV2 架构进行训练的，该数据集包括 91 个类别（包括背景类）。网络的配置细节可以在此处找到。

这个示例是如何工作的？

SSD 网络通过网络的单次前向传播执行对象检测和定位任务。TensorFlow SSD 网络是使用 MSCOCO 数据集在 InceptionV2 架构上训练的。

该示例使用 TensorRT 插件来运行 SSD 网络。为了使用这些插件，需要预处理 TensorFlow 图，并使用 GraphSurgeon 实用程序进行此操作。

该网络的主要组件包括图像预处理器、特征提取器、盒子预测器、网格锚点生成器和后处理器。

图像预处理器
图像预处理图的步骤负责调整图像的大小。图像被调整为 300x300x3 大小的张量。该步骤还执行图像的归一化，使得所有像素值都在范围 [-1, 1] 内。

特征提取器
图的特征提取器部分在经过预处理的图像上运行 InceptionV2 网络。生成的特征图被用于锚点生成步骤，用于为每个特征图生成默认的边界框。

在此网络中，用于锚点生成的特征图的大小为 [(19x19), (10x10), (5x5), (3x3), (2x2), (1x1)]。

盒子预测器
盒子预测器步骤接受高级别特征图作为输入，并为每个特征图的每个编码盒子生成一个盒子编码（x-y 坐标）列表以及每个编码盒子的类别分数列表。然后将这些信息传递给后处理器。

网格锚点生成器
此步骤的目标是为每个特征图单元格生成一组默认边界框（根据配置中提到的比例和长宽比）。这实现为 TensorRT 中的一个名为 gridAnchorGenerator 插件的层。注册的插件名称是 GridAnchor_TRT。

后处理器
后处理器步骤执行生成网络输出的最后步骤。所有特征图的边界框数据和置信度分数都将与预先生成的默认边界框（在 GridAnchorGenerator 命名空间中生成）一起输入到此步骤中。然后执行 NMS（非最大抑制），基于置信度阈值和 IoU（交并比）重叠来删掉大部分边界框，仅保留每类的前 N 个边界框。这实现为 TensorRT 中的名为 NMS 的插件层。注册的插件名称是 NMS_TRT。

注意： 本示例还实现了另一个名为 FlattenConcat 的插件，用于将输入扁平化然后连接结果。在输入传递给后处理器之前，会将位置和置信度数据应用于这个插件，因为 NMS 插件需要数据采用这种格式。

关于插件如何实现的详细信息，请参见 tensorrt/samples/sampleUffSSD 目录中的 sampleUffSSD.cpp 文件中的 FlattenConcat 插件和 FlattenConcatPluginCreator 的实现。

具体来说，此示例执行以下步骤：

处理输入图
准备数据
sampleUffSSD 插件
验证输出

处理输入图

TensorFlow SSD 图具有一些目前在 TensorRT 中不受支持的操作。通过对图进行预处理，我们可以将图中的多个操作合并为单个自定义操作，这些操作可以在 TensorRT 中实现为插件层。目前，预处理器提供了将图中的所有节点（在删除断言后没有输出的节点）拼接成单个自定义节点的能力。

要使用预处理器，应该使用 convert-to-uff 实用程序并带有配置文件的 -p 标志来调用它。配置脚本还应包含将嵌入到生成的 .uff 文件中的所有自定义插件的属性。目前用于 SSD 的示例脚本位于 /usr/src/tensorrt/samples/sampleUffSSD/config.py。

使用图的预处理器，我们能够删除图中的 Preprocessor 命名空间，将 GridAnchorGenerator 命名空间拼接在一起以创建 GridAnchorGenerator 插件，将 postprocessor 命名空间拼接在一起以获得 NMS 插件，并将 BoxPredictor 中的 concat 操作标记为 FlattenConcat 插件。

TensorFlow 图中有一些操作，如 Assert 和 Identity，在推理时可以移除。Assert 等操作已被移除，剩下的节点（在删除断言后没有输出的节点）将被递归地删除。

Identity 操作将被删除，并且输入将被转发到所有连接的输出。有关图预处理器的附加文档，可以在TensorRT API中找到。

准备数据

生成的网络具有名为 Input 的输入节点，输出节点的名称由 UFF 转换器命名为 MarkOutput_0。

parser->registerInput("Input", DimsCHW(3, 300, 300), UffInputOrder::kNCHW);
parser->registerOutput("MarkOutput_0");

本示例中 SSD 网络的输入是 3 通道 300x300 图像。在示例中，我们对图像进行归一化，使像素值位于范围 [-1,1]。

由于 TensorRT 不依赖于任何计算机视觉库，图像以每个像素的二进制 R、G 和 B 值表示。格式是可移植图像映射（Portable PixMap，PPM），这是一种 netpbm 颜色图像格式。在此格式中，每个像素的 R、G 和 B 值由整数字节（0-255）表示，并按像素存储在一起。

有一个名为 readPPMFile 的简单 PPM 读取函数。

sampleUffSSD 插件

有关如何创建 TensorRT 插件的详细信息，请参见使用自定义层扩展 TensorRT。

用于 convert-to-uff 命令的 config.py 定义应通过修改 op 字段将自定义层映射到 TensorRT 中的插件名称。插件参数的名称也应与 TensorRT 插件所期望的名称和类型完全匹配。例如，对于 GridAnchor 插件，config.py 应该如下所示：

PriorBox = gs.create_plugin_node(name="GridAnchor",
op="GridAnchor_TRT",
numLayers=6,
minSize=0.2,
maxSize=0.95,
aspectRatios=[1.0, 2.0, 0.5, 3.0, 0.33],
variance=[0.1,0.1,0.2,0.2],
featureMapShapes=[19, 10, 5, 3, 2, 1])

这里，GridAnchor_TRT 与已注册的插件名称匹配，参数的名称和类型与插件所期望的名称相同。

如果 config.py 定义如上，NvUffParser 将能够解析网络并使用正确的参数调用适当的插件。

以下是TensorRT中为SSD实现的一些插件层的详细信息。

GridAnchorGeneration 插件
这个插件层实现了TensorFlow SSD网络中的网格锚点生成步骤。对于每个特征图，我们计算每个网格单元的边界框。在这个网络中，有6个特征图，每个网格单元的边界框数量如下：

[19x19] 特征图：3个边界框（19x19x3x4(坐标/边界框)）
[10x10] 特征图：6个边界框（10x10x6x4）
[5x5] 特征图：6个边界框（5x5x6x4）
[3x3] 特征图：6个边界框（3x3x6x4）
[2x2] 特征图：6个边界框（2x2x6x4）
[1x1] 特征图：6个边界框（1x1x6x4）

NMS 插件
NMS 插件生成基于BoxPredictor生成的位置和置信度预测的检测输出。此层有三个输入张量，对应于位置数据(locData)、置信度数据(confData)和先前框数据(priorData)。

检测输出插件的输入必须被扁平化并连接在所有特征图上。我们使用样本中实现的FlattenConcat插件来实现这一点。BoxPredictor生成的位置数据的尺寸如下：

19x19x12 -> 重塑 -> 1083x4 -> 扁平化 -> 4332x1
10x10x24 -> 重塑 -> 600x4 -> 扁平化 -> 2400x1

等等，对于其余特征图也是如此。

连接后，locData 输入的维度约为7668x1。

BoxPredictor生成的置信度数据的尺寸如下：

19x19x273 -> 重塑 -> 1083x91 -> 扁平化 -> 98553x1
10x10x546 -> 重塑 -> 600x91 -> 扁平化 -> 54600x1

等等，对于其余特征图也是如此。

连接后，confData 输入的维度为174447x1。

Grid Anchor Generator 插件生成的先前数据有6个输出，它们的维度如下：

输出 1 对应于19x19特征图，维度为2x4332x1
输出 2 对应于10x10特征图，维度为2x2400x1

等等，对于其他特征图也是如此。

注意： 输出中有两个通道，因为一个通道用于存储每个坐标的方差，这在NMS步骤中使用。连接后，priorData 输入的维度约为2x7668x1。

struct DetectionOutputParameters
{
bool shareLocation, varianceEncodedInTarget;
int backgroundLabelId, numClasses, topK, keepTopK;
float confidenceThreshold, nmsThreshold;
CodeTypeSSD codeType;
int inputOrder[3];
bool confSigmoid;
bool isNormalized;
};

shareLocation 和 varianceEncodedInTarget 用于Caffe SSD网络实现，所以对于TensorFlow网络，它们应该分别设置为 true 和 false。confSigmoid 和 isNormalized 参数对于TensorFlow实现是必要的。如果将 confSigmoid 设置为 true，它将计算所有置信度得分的sigmoid值。isNormalized 标志指定数据是否被规范化，对于TensorFlow图，它被设置为 true。

验证输出

在创建生成器之后（参见在C++中构建引擎）并序列化引擎之后（参见在C++中序列化模型），我们可以执行推理。关于反序列化和运行推理的步骤在在C++中执行推理中进行了概述。

SSD网络的输出是人类可解释的。最后的NMS等后处理工作在NMS插件中完成。结果被组织为元组，每个元组有7个元素，分别是图像ID，对象标签，置信度分数，边界框左下角的(x,y)坐标，以及边界框的右上角的(x,y)坐标。可以使用writePPMFileWithBBox函数在输出PPM图像上绘制这些信息。visualizeThreshold参数可用于控制图像中对象的可视化。目前它被设置为0.5，所以输出将显示所有置信度为50%及以上的对象。

TensorRT API层和操作

在这个示例中，使用了以下层。有关这些层的更多信息，请参阅TensorRT开发者指南：层文档。

激活层
激活层实现了逐元素激活函数。具体而言，本示例使用了类型为kRELU的激活层。

连接层
连接层沿着通道维度连接多个非通道大小相同的张量。

卷积层
卷积层计算带有或不带有偏差的2D（通道、高度和宽度）卷积。

填充层
填充层在最内部的两个维度上实现了零填充。

插件层
插件层是用户定义的，提供了扩展TensorRT功能的能力。有关更多详细信息，请参阅通过自定义层扩展TensorRT。

池化层
池化层在通道内进行池化。支持的池化类型有maximum、average和maximum-average blend。

尺度层
尺度层实现了每个张量、每个通道或每个元素的仿射变换和/或常数值的幂运算。

Shuffle层
Shuffle层为张量实现了重塑和转置运算符。

先决条件

安装 UFF 工具包和图形外科医生；根据您的 TensorRT 安装方法，选择您所使用的方法以安装工具包和图形外科医生，可以参考 TensorRT 安装指南：安装 TensorRT 获取详细说明。
下载 ssd_inception_v2_coco TensorFlow 训练模型。
使用 UFF 转换器对 TensorFlow 模型进行预处理。
1. 从之前步骤下载的目录中将 TensorFlow 协议缓冲文件（frozen_inference_graph.pb）复制到工作目录中（例如 /usr/src/tensorrt/samples/sampleUffSSD/）。
2. 运行以下命令进行转换。
```
convert-to-uff frozen_inference_graph.pb -O NMS -p config.py
```
```
 这将保存转换后的 `.uff` 文件到与输入相同的目录，并命名为 `frozen_inference_graph.pb.uff`。`config.py` 脚本指定了 SSD TensorFlow 图所需的预处理操作。`config.py` 脚本中使用的插件节点和插件参数应与 TensorRT 中的注册插件匹配。
```
1. 将转换后的 .uff 文件复制到数据目录，并将其重命名为 sample_ssd_relu6.uff <TensorRT 安装目录>/data/ssd/sample_ssd_relu6.uff。
该示例还需要一个包含用于训练模型的所有标签的 labels.txt 文件。该网络的标签文件是 <TensorRT 安装目录>/data/ssd/ssd_coco_labels.txt。

运行示例

在 <TensorRT 根目录>/samples/sampleUffSSD 目录中运行 make 来编译该示例。该二进制文件将被创建在 <TensorRT 根目录>/bin 目录中。
```
cd <TensorRT 根目录>/samples/sampleUffSSD
make
```
其中 <TensorRT 根目录> 是您安装 TensorRT 的位置。
运行示例以执行对象检测和定位。

要在 FP32 模式下运行示例：
```
./sample_uff_ssd
```
要在 INT8 模式下运行示例：
```
./sample_uff_ssd --int8
```
注意： 要在 INT8 模式下运行网络，请参考 BatchStreamPPM.h，了解如何执行校准。目前，我们需要一个名为 list.txt 的文件，其中列出了位于 <TensorRT 安装目录>/data/ssd/ 文件夹中用于校准的所有 PPM 图像。用于校准的 PPM 图像也可以位于同一文件夹中。

验证示例是否成功运行。如果示例成功运行，您应该会看到类似以下的输出：

&&&& RUNNING TensorRT.sample_uff_ssd # ./build/x86_64-linux/sample_uff_ssd
[I] ../data/samples/ssd/sample_ssd_relu6.uff
[I] 开始解析模型...
[I] 完成解析模型...
[I] 开始构建引擎...
I] 批处理数量 1
[I] 数据大小 270000
[I] *** 反序列化
[I] 推断所用时间为 4.24733 毫秒。
[I] 保留数量 100
[I] 在图像 0（../../data/samples/ssd/dog.ppm）中检测到狗，置信度为 89.001，坐标为 (81.7568, 23.1155)，(295.041, 298.62)。
[I] 结果保存在 dog-0.890010.ppm 中。
[I] 在图像 0（../../data/samples/ssd/dog.ppm）中检测到狗，置信度为 88.0681，坐标为 (1.39267, 0)，(118.431, 237.262)。
[I] 结果保存在 dog-0.880681.ppm 中。
&&&& PASSED TensorRT.sample_uff_ssd # ./build/x86_64-linux/sample_uff_ssd

这个输出表明示例成功运行；“PASSED”。

示例 `--help` 选项

要查看可用选项的完整列表及其说明，请使用 -h 或 --help 命令行选项。

附加资源

以下资源提供了有关 TensorFlow SSD 网络结构的更深入理解：

模型

TensorFlow 检测模型库

网络

ssd_inception_v2_coco_2017_11_17

数据集

MSCOCO 数据集

文档

NVIDIA TensorRT 示例简介
使用 C++ API 使用 TensorRT
NVIDIA TensorRT 文档库

许可协议

有关使用、复制和分发的条款和条件，请参阅TensorRT 软件许可协议文档。

更新日志

2019年3月
重新创建、更新和审阅此 README.md 文件。

已知问题

在 INT8 模式下运行网络时，可能会存在一些精度损失，导致某些对象无法被检测到。一般观察是，>500 张图像是用于校准的良好数量。
在 Windows 上，Python 脚本 convert-to-uff 不可用。您可以在 Linux 机器上生成所需的 .uff 文件，然后复制到 Windows 上以运行此示例。