基于NVIDIA Jetson TX2的道路场景分割①

李诗菁， 卿粼波， 何小海， 韩 杰

(四川大学 电子信息学院， 成都 610065)

近年， 深度学习的迅猛崛起给各个科技领域的发展带来了巨大影响. 在计算机视觉领域， 图像语义分割是重要的基础研究问题之一. 传统的图像分割一般采用基于阈值、边缘、区域的分割方法， 根据颜色纹理等人工标定的特征完成分割， 过程复杂且局限性较大[1]. 基于深度学习的图像语义分割网络的分割性能远超传统算法， 可应用到诸多领域如无人驾驶、智能安防、交通监控、机器人中[2]. 在较有代表性的图像分割网络FCN (Fully Convolutional Networks)[3]、SegNet(SegNet: A Deep ConvolutionalEncoder-Decoder Architecture for ImageSegmentation)[4]、PSPNet(Pyramid Scene Parsing Network)[5]和ENet(ENet: A Deep Neural Network Architecture forReal-Time Semantic Segmentation)中， FCN是最早的图像分割网络之一， 准确度较低且耗时长； PSPNet分割的像素精确度高， 但耗时较长； ENet分割速度较快， 但分割效果较不理想； SegNet兼具较高的分割准确度与速度， 是完成道路场景分割任务较好的选择.

随着人工智能的不断发展， 深度学习网络在嵌入式平台上部署的需求也日益增长. 目前国内外推出的可部署深度学习的主流芯片一般分为FPGA芯片和GPU芯片两类. 在FPGA上已实现了一些小规模的神经网络， 如Zhang等人在Xilinx Virtex 7485T FPGA上实现了AlexNet[6]， 卢冶、陈瑶等人在XC7Z020 FPGA上实现了LeNet-5、Cifarnet网络[7]. 基于FPGA的深度学习应用的功耗小于GPU， 但受限于FPGA开发难度与硬件资源， 目前大规模复杂结构的网络在嵌入式上的部署还是更多地基于GPU实现. 已有许多深度学习网络逐渐被应用到了嵌入式ARM+GPU平台如NVIDIA JetsonTX2 (简称TX2)上， 构建了目标检测、感知导航等系统[8，9]. 因分割网络相较于图像分类、目标检测更难实现， 大部分高准确度的分割网络在低功耗的嵌入GPU平台上运行速度极慢， 而较少地被应用到嵌入式平台上.

因此， 针对无人驾驶与交通监控等应用背景， 本文选取SegNet网络在嵌入式平台TX2上实现道路场景理解， 并采用模型简化与转换网络模型为加速引擎的方式， 在基本保持原网络分割准确度的情况下大幅度提升了分割网络在嵌入式平台中的推理速度. 本文首先合并BN层简化了网络模型， 然后基于英伟达的TensorRT2前向推理加速器采用水平层集成， 关联层消除， 权值精度减半等优化措施， 将网络模型转换成了加速引擎进行前向推理. 下面将分别就网络结构与模型简化、基于TensorRT2的加速引擎构建两部分进行介绍.

1 网络结构与网络模型简化

1.1 网络结构

本文采用SegNet网络， 实现将道路场景分为行人、道路、天空、标志、车辆、交通标线、杆状物、自行车、侧路、植物、建筑、围墙共12类的任务.

SegNet网络基本结构为自动编-解码器结构， 采用VGG16前13层卷积层作为编码器， 后接13层解码网络与一个分类层. SegNet关键在于存储了编码网络中每个的池化层中的最大值与其空间信息， 用于对应解码器的非线性上采样， 这极大精确了分割中的边界定位， 减少了编码器到解码器的参数量， 使得SegNet在速度与内存利用上都具有很大优势.

1.2 网络模型简化

在将网络模型转换为加速引擎之前， 本文对SegNet网络进行了模型简化. SegNet网络中引入了BN层， 这种层可在训练时期加速网络收敛， 但在推理过程中会增大内存消耗， 降低推理速度. 可将BN层与其相连的卷积层合并以简化模型加速推理， 如图1所示.

图1 BN层合并

卷积层与BN层都是线性变换， 提供了合并的可能性. 设输入的mini-batch: {x1··n}， BN层变化如式(1):

式(1)中，i∈[0,n]， 且i为整数与是训练中得到的参数为趋于0的常数.

2 基于TensorRT2的加速引擎构建

SegNet在TX2上的推理速度极慢， 因此本文采用NVIDIA推出的TensorRT2对网络模型进行加速.

2.1 基本流程

TensorRT是NVIDIA公司推出的深度学习网络推理引擎， 可优化已有网络模型， 大幅提升神经网络在如机器人、自动驾驶平台上的推理速度. 目前NVIDIA公司已推出四个版本的TensorRT. 不同版本TensorRT可支持的深度学习框架不同.

因为本文使用的SegNet网络采用caffe架构， 且需要插入TensorRT不支持的层， TensorRT2版本已经能满足本文需求. 故本文使用TensorRT2， 基于训练后的网络模型生成加速推理引擎， 完成高效的GPU推理性能优化. 加速后的模型无需深度学习架构支持， 对平台的依赖性极小， 无论在高性能GPU(如NVIDIA Titan X)上， 还是在嵌入式平台(例如NVIDIA TX2)上都拥有不俗的加速能力. 加速引擎构建的基本流程如图2.

图2 加速引擎构建流程

构建过程中首先创建网络， 导入网络模型并利用NvcaffeParser进行解析， 然后根据网络结构与定义采取优化措施生成加速引擎. 生成的加速引擎可序列化存储到磁盘中. 在执行阶段反序列化加速引擎， 进行输入与输出口的绑定、GPU显存分配与计算内核启动，执行加速引擎得到分割结果.

2.2 自定义层添加

在上文的网络模型解析过程中， 需将原网络模型中的层转换为其支持的形式. 加速引擎中除部分常见层如卷积层、激活层、全连接层等可直接转换外， 其它TensorRT不支持的层需自行定义， 下文将这些层统称为自定义层. TensorRT中支持的层可直接由上文提到的NvcaffeParser进行解析， 自定义层则需采用plugin接口添加入NvcaffeParser中.

自定义层plugin添加的主要流程如下: 输出确定、层配置、工作空间分配、资源管理、序列化与层执行. 在加速引擎中， 一个层的输入或输出定义为tensor， tensor具有数据类型与三个维度分量， 即通道数C、宽度W与高度H. 在开始插入层时， 需对插入层的输出进行定义， 需确定插入层的输出数目与输出tensor的三个维度分量{C，H，W}； 完成输出确定后， 将进行层配置， 该过程主要获取输入tensor的形式； 在层配置后， 加速引擎会分配临时的工作空间在自定义层之间共享以达到内存利用率最大化； 同时还需要进行资源管理的配置， 主要是通过层资源初始化与销毁来完成资源分配与释放； 在进行完成上述部分后， 需判断加速引擎是否需要序列化到磁盘中， 若需序列化， 则将自定义层的参数与网络的其余部分合并以便后续整个网络的序列化存储； 最后在层执行阶段， 主要完成层的算法实现， 如果未选择序列化存储， 则直接在资源分配后执行该过程， 若已进行序列化存储， 则提取序列化参数后执行该过程， 整体过程如图3所示.

图3 plugin层插入流程

本文将网络中需两个输出口的pooling层、upsample层与argmax分类层添加到加速模型中， 这些添加的层是网络的关键结构， 如图4所示. 本文将pooling层分出了两个输出口， 将池化结果送入下一级编码器的同时将池化的最大特征值与位置信息输入了对应upsample层中. upsample层将接收的池化最大特征值按照位置信息填入上采样的稀疏特征图中， 空缺位置补零， 作为解码器的输入， 如图5. 最终经过层层解码后的特征图输入argmax层， 得到每个像素最大可能性的分类结果.

图4 plugin层结构

图5 plugin层实现过程

2.3 优化措施

本文的优化措施主要分为图1中的引擎生成阶段的优化与图3自定义层添加中执行阶段enqueue( )时的CUDA编程优化.

网络模型中的各个层经过解析之后， 在引擎生成时主要采用了以下措施进行优化: 清除未使用的输出层以避免不必要的计算； 整合垂直结构的卷积层、ReLU激活层、Bias偏置项为CBR层； 将拥有同一输入、相同操作、相似参数和不同输出的水平层集成为一个层，例如图6； 在平台支持的情况下， 将数据精度需求从32位float降低为半精度的fp16， 以此提升计算效率.

在自定义层添加时执行层算法的函数enqueue( )中， 本文采用了CUDA编程的方式将算法中大量耗时的计算分配到GPU上并行实现. 原网络最后的分类层argmax层为CPU实现， 若同样采用CPU实现， 需将输入层的数据从GPU上复制到CPU上. 数据复制将耗费大量时间， 且argmax层算法在CPU上的执行时间也过长. 因此本文在GPU上实现了argmax层前馈算法. 在CUDA C中的host函数只可由CPU调用， device函数则只可运行在GPU上， 且不能调用常用函数. 因此argmax层前馈算法中调用的部分算法， 并不能在GPU上调用. 为此本文针对网络的实际应用简化了相应算法以便GPU实现. 最终在得到每个像素值12类分类结果后， 在一维索引空间中并行360*480个线程，将位置一一对应的特征值的12个分类结果进行排序，得到整幅图像像素级别的分类.

图6 水平层集成

3 实验测试

3.1 网络训练

本文采用GPU模式进行网络训练， 硬件为基于Pascal架构的NVIDIA Titan X显卡， 该卡显存为12 GB.

为满足针对交通监控与自动驾驶的道路场景理解的目的， 本文从包含城市交通场景的数据库KITTI[10]、Camvid[11]与Cityscapes[12]中选取了1300张图像作为训练数据集， 并将原图与标签裁剪为统一尺寸. 此外， 选取的网络训练需要单通道标签， 故将原数据集中RGB标签转换为单通道， 并将每个像素转换为代表天空、行人、围墙、植物、交通标志、侧道、道路、车辆等12个分类的灰度值. 在训练时， 为提高分类准确度， 使用预先训练好的VGG16网络模型的权值对SegNet编码网络的权值进行初始化. 在20 000次iteration后， 网络准确度就已经达到80%以上， 在迭代36 000次时，达最高准确度92%， 迭代38 000次时达到最低损失0.019.

3.2 测试结果

在训练得到最优网络模型的基础上， 本文完成了加速引擎(下文称为SegEng)的构建， 并将其序列化后存储到磁盘上. NVIDIA TX2平台支持fp16， 故本文生成SegEng时采用了半精度模式， 将其大小缩减为SegNet模型的一半. 在 NVIDIA TX2 上， 本文对 GPU模式下使用cudnn6.0加速的SegNet与SegEng的前向推理性能进行了测试对比.

图7 分割结果

分割图对比如图7所示. 在图像分割领域， 常用以下三个指标衡量分割精度: 像素精确度(G)表示图像中被正确分类的像素比； 类平均精确度(C)表示所有类被正确分类的平均值； 交除并均值(Mean Intersection over Union， MIoU)代表所有类的IoU平均值.

本文从测试集中选取图片， 使用上述三个指标对迭代了40 000次得到的SegNet模型与SegEng的分割结果进行评估， 得到表1.

表1 分割精度对比(单位: %)

原始模型中的权值精度为32位， 加速引擎舍弃了权值中后16位的冗余信息， 将权值精度降低为16位，权值中丢弃的后16位信息对最终的像素级别分类概率影响极小， 由表1中可看到并未影响到最终的分割准确度.

在NVIDIA TX2上， 本文对不同分辨率的1000帧视频进行了10次测试后取平均值， 得到了表2中SegNet与SegEng推理时间的对比.

表2 单帧图像推理时间对比(单位: ms)

在经过上文的层合并、网络层优化、CUDA加速与TensorRT内部的卷积优化、内存优化、精度减半等优化措施后， 由表 2可看出， 在 TX2上， 相较于SegNet， 本文生成的SegEng达到了约十倍的加速比，大幅度提高了SegNet在嵌入式平台上的推理速度. 在分割 480*320大小的图片时， 分割速度可达 10 fps， 具有在实际道路场景解析的应用潜力.

4 结语

本文针对无人驾驶与交通监控的应用背景， 采用了图像语义分割网络SegNet完成将道路场景分割为行人、车辆、道路、植物、侧道、交通标志、自行车等12类对象的任务. 同时本文为提升分割网络在嵌入式平台上的推理速度， 对网络模型进行了简化， 然后基于NVIDIA推出的TensorRT2， 采用集成水平层、消除多余输出层、采用权值精度减半、CUDA并行编程等优化措施， 完成了网络模型加速并生成了加速引擎.在NVIDIA TX2嵌入式平台上， 本文生成的加速引擎无需深度学习架构支持， 在分割精度基本无影响情况下， 推理速度可达原网络的十倍， 为复杂结构分割网络在嵌入式平台上的应用提供了支持.