【参考算法】地平线双目深度估计参考算法StereoNetPlus-v1.2.1

0 概述

1 性能精度指标

StereoNetPlus模型配置：

性能精度表现：

EPE（End-Point Error）是以像素为单位的平均视差误差。

2 模型介绍

• Backbone：使用共享权重的MixVarGENet来提取双目图像的多尺度特征；

• Neck：在原图1/8、1/16、/1/32特征级别使用FPN网络，提取每个尺度的特征，然后用以构建cost volume；

• Head：使用内积计算左右特征图的相关性，构建左右图像视差cost volume，然后做Adaptive intra-scale aggregation和Adaptive cross-scale aggregation，最后预测优化视差；另外，将左图像的1/2、1/4和1/8尺度特征融合并做上采样，用于视差细化；

• Post Process：根据head中获取的视差和cost volume计算最终的视差。

2.1 模型优化点

在网络结构上，相比于公版实现，地平线做了如下更改：

1. Backbone： 将公版的孪生网络替换为地平线进行针对性优化的MixVarGENet，获得了性能和精度的提升；

2. cost volume构建与融合：采用了AANet基于相关性的多尺度cost volume构建方式，并将1/8， 1/16， 1/32原图尺度下（AANet为1/4， 1/8， 1/16尺度）的多层correlation cost volume进行融合，最终使用融合后的1/8视差做refinement，实现了精度和性能的最佳平衡；
3. 初略视差到精细视差：采用了CoEx的思路，将左图进行超像素上采样，从而获得原始分辨率的视差图，然后与融合后的cost volume联合预测视差；

4. 将后处理中的unfold操作优化为conv进行计算，实现了在BPU上的加速。

2.2 源码说明

2.2.1 Config文件

configs/disparity_pred/stereonet/stereonetplus_mixvargenet_sceneflow.py为 stereonetplus的配置文件，定义了模型结构、数据集加载，和整套训练流程，所需参数的说明在算子定义中会给出。配置文件主要内容包括：

注： 如果需要复现精度，config中的训练策略最好不要修改。否则可能会有意外的训练情况出现。

2.2.2 Backbone

StereoNetPlus参考算法将backbone由公版的共享权重的孪生神经网络替换为了地平线进行针对性优化的MixVarGENet。MixVarGEBlock包括head op， stack op，downsample layers，fusion layers四个基本模块，其中后三个模块都是可选模块。如下为MixVarGEBlock的结构图：

并且，head_op 和stack_op都是由BasicMixVarGEBlock（如mixvarge_f2，mixvarge_f4，mixvarge_f2_gb16）这样的基本单元构成。

2.2.3 Neck

StereoNetPlus的Neck网络采用特征金字塔网络FPN(Feature Pyramid Networks)，它采用自顶向下的层次结构来提取多尺度的高层语义特征。FPN网络结构如图所示：

2.2.4 Head

2.2.4.1 cost volume构建

cost volume用于在双目匹配中衡量左右图像像素匹配的程度，StereoNetPlus通过计算左右视图的相关性来构建cost volume，计算公式如下：

相关代码如下：

在计算cost volume之前，判断i是否为正数。i为正时需要对cost_volume做padding。

2.2.4.2 Adaptive Aggregation

构建cost volume后，然后用几个堆叠的自适应聚合模块来聚合原始cost volume，其中自适应聚合模块由3个自适应尺度内聚合（Intra-Scale Aggregation ， ISA）模块和用于3个金字塔级别的自适应跨尺度聚合（Cross-Scale Aggregation ，CSA）模块组成。ISA和CSA的示意图如下所示：

ISA

在视差非连续时，边缘位置总会有一圈连续的错误匹配值，为了缓解这种edge-fattening问题，使用3个残差模块 BasicResBlock对每个尺度的cost volume进行聚合，BasicResBlock的结构图如下所示：

相关代码：

CSA

双目图像进行下采样后，在相同的patch尺寸下，纹理信息将更具鉴别性，所以跨尺度成本聚合算法中引入了多尺度交互。最终的cost volume是通过对不同尺度的成本聚合结果进行自适应组合得到的，公式如下：

相关代码如下:

ISA和CSA的forward函数为：

代码路径：

/usr/local/lib/dist-packages/hat/models/task_modules/stereonet/headplus.py

预测视差

获取到多尺度融合的Cost volume后，首先使用softmax函数将其转化为视差的分布；然后使用Conv代替Unfold操作进行张量的切分，这样就获得了视差预测，相关操作如下所示：

> 上图中nn.ZeroPad2d()函数的作用是为了保证Conv2x2计算后的shape为68x120，即1/8原图尺度。

相关代码：

使用Conv代替Unfold操作的代码为：

代码路径：

/usr/local/lib/dist-packages/hat/models/task_modules/stereonet/headplus.py

2.2.4.3 超像素上采样

在上采样分辨率中的每个像素处的最终视差估计是通过其周围的“超像素”的加权平均值获得的。StereoNetPlus参考算法借鉴了CoEx的做法，将1/2、1/4、1/8分辨率下的左图特征进行上采样获得原始图像分辨率的视差图。

如下为超像素上采样的网络结构示意图，首先使用DeconvResModule融合金字塔特征P1/2、P1/4、P1/8的左边特征，DeconvResModule的结构图如右侧所示，它首先使用Convtranspose将小尺寸特征图input1上采样至input2大小，然后将二者的特征进行相加、Conv等操作；最后使用ConvTranspose将融合的特征图上采样至原图大小544x960。如下为流程图：

代码路径：

/usr/local/lib/dist-packages/hat/models/task_modules/stereonet/headplus.py

2.2.5 Post Process

获取到原图尺寸视差估计后，基于cost volume中的视差分布信息获取最终的视差预测。首先，使用最近邻插值将1/8原图分辨率下的cost volume上采样至原图分辨率下（544x960）；然后将其与上采样后的超像素对应点相乘并求和，这样就获得了每个像素处的最终视差估计；最后，乘以最大视差192就获得最终的视差，相关流程如下图所示：

相关代码：

代码路径：

/usr/local/lib/dist-packages/hat/models/task_modules/stereonet/post_process.py

3 浮点模型训练

3.1 Before Start

3.1.1 环境部署

Stereonetplus示例位于OE包中的ddk/samples/ai_toolchain/horizon_model_train_sample下，其结构为：

拉取docker环境

权重获取路径见：horizon_model_train_sample/scripts/configs/disparity_pred/stereonet/README.md

3.1.2 数据集下载

在开始训练模型之前，需要准备好数据集。首先下载 SceneFlow 数据集 ， 然后在此处下载训练数据和验证数据集对应的文件列表 SceneFlow_finalpass_train.txt 和 SceneFlow_finalpass_test.txt 。

下载后，解压并按照如下方式组织文件夹结构：

3.1.3 数据集打包

为了提升训练的速度，需要对数据信息文件进行打包，将其转换成lmdb格式的数据集。只需要运行下面的脚本，就可以成功实现格式转换：

src-data-dir为解压后的SceneFlow数据集目录，即解压后的数据集目录SceneFlow；

target-data-dir为打包后数据集的存储目录；

num-worker为执行线程数

数据集打包命令执行完毕后会在target-data-dir下生成train_lmdb和test_lmdb:

3.1.4 config配置

在进行模型训练和验证之前，需要对configs文件中的部分参数进行配置，一般情况下，我们需要配置以下参数：

• device_ids、batch_size_per_gpu：根据实际硬件配置进行device_ids和每个gpu的batchsize的配置；

• ckpt_dir:float、calibration、qat的权重路径配置，权重下载链接在config文件夹下的README中；

• data_loader中的data_path：2.1.3节打包的train_lmdb数据集路径；

• val_data_loader中的data_path：2.1.3节打包的test_lmdb数据集路径；

• infer_cfg：指定模型输入，在infer.py脚本使用时需配置；

3.2 浮点模型训练

在configs/disparity_pred/stereonet/stereonetplus.py下配置参数，需要将相关硬件配置device_ids和权重路径ckpt_dir数据集路径data_rootdir配置修改后使用以下命令训练浮点模型：

3.3 浮点模型精度验证

通过指定训好的float_checkpoint_path，使用以下命令验证已经训练好的模型精度:

4. 模型量化和编译

4.1 Calibration

为加速QAT训练收敛和获得最优量化精度，建议在QAT之前做calibration，其过程为通过batchsize个样本初始化量化参数，为QAT的量化训练提供一个更好的初始化参数，和浮点训练的方式一样，将checkpoint_path指定为训好的浮点权重路径。

通过运行下面的脚本就可以开启模型的Calibration：

4.2 Calibration 模型精度验证

calibration完成以后，可以使用以下命令验证经过calib后模型的精度:

验证完成后，会在终端输出calib模型在验证集上的EPE精度。

4.3 量化模型训练

Calibration完成后，就可以加载calib权重开启模型的量化训练。 量化训练其实是在浮点训练基础上的finetue，具体配置信息在config的qat_trainer中定义。量化训练的时候，初始学习率设置为浮点训练的十分之一，训练的epoch次数也大大减少。和浮点训练的方式一样，将checkpoint_path指定为训好的calibration权重路径。

通过运行下面的脚本就可以开启模型的qat训练：

4.4 qat模型精度验证

量化训练完成后，通过运行以下命令验证qat模型的精度:

4.5 量化模型精度验证

通过运行以下命令验证量化模型的精度:

4.6 仿真上板精度验证

除了上述模型验证之外，我们还提供和上板完全一致的精度验证方法，可以通过下面的方式完成：

4.6 量化模型编译

opt为优化等级，取值范围为0~3，数字越大优化等级越高，运行时间越长。

运行后，out-dir的目录下会产出以下文件：

5. 其他工具

5.1 结果可视化

如果您希望查看定点模型通过双目图像估计视差和深度的效果，tools文件夹下面提供了预测及可视化的脚本infer.py，参考命令如下：

参数说明：

1. 输入图像的路径通过config文件中infer_cfg字段中的infer_inputs配置；

2. --save-path：可视化结果保存路径;

可视化结果：

1. 可视化结果从左至右为：left左图、right右图、disparity视差图、根据视差计算的depth深度图；

2. 视差图中的颜色代表：颜色越红，视差值越小；颜色越蓝，视差值越大；

3. 深度图中的颜色代表：颜色越红，对应的像素深度值则越小；颜色越蓝，对应的像素深度值则越大；

4. 视差图到深度图的计算原理可参考【双目深度估计】—原理理解中的几何法，相应代码可以参考config文件中的process_outputs函数。

5.2 导出onnx

如果需要导出onnx模型，tools文件夹下面提供了onnx导出脚本export_onnx.py，参考命令如下：

导出的onnx模型为量化训练后的QAT模型，保存在模型权重路径ckpt_dir目录下。

6 板端部署

本节将介绍hbm模型编译完成后，在板端使用dnn工具进行性能评测和运行AI-Benchmark示例进行性能和精度评测的流程。

6.1 上板性能实测

模型编译成功后，可以在板端使用hrt_model_exec perf工具评测hbm模型的FPS，参考命令如下：

命令运行结束后，会在本地会产出profile.log和profile.csv日志文件，用以分析算子耗时和调度耗时。

6.2 AI Benchmark示例

OE开发包中提供了StereoNetPlus的AI Benchmark示例，用以在板端进行性能和精度的评测，ddk/samples/ai_benchmark/j5/qat/script/disparity_pred/stereonet_plus 目录下提供了三个评测脚本：

• fps.sh：实现多线程fps统计（多线程调度，用户可以根据需求自由设置线程数）；

• latency.sh：实现单帧延迟性能统计，包含模型推理和后处理延迟（单线程，单帧）；

• accuracy.sh：用于精度评测。

• AI Benchmark示例需要使用交叉编译工具编译后，然后将文件夹传输到板端运行，建议在docker中进行交叉编译流程。执行ddk/samples/ai_benchmark/code/ 目录下的build_qat_j5.sh脚本即可一键编译真机环境下的可执行程序：

编译完成后，将ai_benchmark/j5/qat文件夹拷贝至板端：

然后在板端执行qat/script/disparity_pred/stereonet_plus文件夹下的评测脚本。进入到需要评测的模型目录下，运行latency.sh 即可测试出单帧延迟, 如下所示：

终端输出的延迟包含模型推理耗时（Infer latency）和后处理耗时(Post process)。

运行fps.sh 即可测试出双核多线程fps, 如下所示：