0 概述
近年来，计算机视觉领域的发展推动了机器感知系统的快速发展。光流估计是计算机视觉领域的基本问题之一，在自动驾驶和机器人领域中有着广泛的应用。具体来说，光流估计有助于车辆感知周围环境的时间连续性，因此它在基于时间序列的任务中发挥了重要作用，如目标跟踪、语义分割和运动分割等；同时光流估计推动了姿态识别和姿态估计等下游应用的发展，有助于机器人感知交互系统的优化。-
光流估计任务为求解连续两帧图像中目标物体或者目标像素点的偏移量，偏移量往往用二维向量表示。按照是否选取图像中稀疏点进行光流估计，可以分为稀疏光流和稠密光流，PWCNet面向稠密光流估计。PWCNet拥有参数量小和轻量化优点的同时，还有着优秀的估计精度，使其成为其中最热门的算法之一，因此地平线集成了光流估计PWCNet算法。本文将详细介绍光流估计PWCNet参考算法的原理和部署。
该示例为参考算法，仅作为在J5上模型部署的设计参考，非量产算法
1 性能精度指标
模型参数：
数据集
Input Shape
Backbone
Head
FlyingChairs
1x6x384x512
PwcNetNeck
PwcNetHead
FlyingChairs数据集参考：FlyingChairs数据集
模型性能精度表现：
模型
性能（FPS/双核）
EndPointError（定点模型）
PWCNet
161.52
1.4075
2 模型介绍
PWCNet网络模型由NVIDIA在“CVPR, 2018”提出，是一个轻量且高效的基于CNN的光流估计模型，相比稠密光流估计的主流模型FlowNet 2.0，其网络大小缩小了17倍并且对于训练更加友好。此外截至论文发稿，PWCNet在“MPI Sintel final pass”和“KITTI 2015 benchmarks”上的性能和精度优于发布的所有光流估计算法，成为光流估计任务的SOTA（State-Of-The-Art）模型，同时是基于深度学习的光流估计中非常基础且具有重要意义的一个网络模型。
-
PWCNet的模型设计主要基于三个简单高效且成熟运用的原则：金字塔多尺度特征提取（Pyramidal Processing）、光流映射（Warping）和相关性匹配成本量（Cost Volume）。其模型构成包括多尺度特征金字塔、光流映射层、相关性匹配成本层、光流估计器和Context Network：
- **多尺度特征金字塔：**不同于直接使用图像金字塔，特征金字塔对阴影和光照变化更为鲁棒，PWCNet基于多尺度特征金字塔的设计使其对不同尺度的光流（不同大小的像素点运动偏移量）拥有更准确的感知，因为对于一个LLL层的特征图金字塔，顶层特征图一个像素点的位移映射回原图则是2L−12^{L-1}2L−1个像素点的位移。此外用多尺度特征金字塔替换传统的子网络串联，极大地减少可学习参数量的同时，减小了模型训练的复杂度；
- **光流映射层+相关性匹配成本层：**光流映射层实现了将上层输出的光流（除金字塔顶层）应用到当前层特征图并生成映射后的特征图，相关性匹配成本层衡量相同尺度的第一帧图像特征图和第二帧图像映射后的特征图之间的像素相关性。光流映射层和相关性匹配层共同使每一层的光流（除金字塔顶层）都在上一层的基础上进一步Refine，实现了对输出光流的Coarse-to-fine。此外对于光流的非平移运动，光流映射操作对几何扭曲有补偿作用且可以将图像块恢复到正确比例。光流映射层和相关性匹配层没有可学习的参数，实现了模型参数量的减少和轻量化的目的；
- **光流估计器：**光流估计器由CNN网络和反卷积层组成，CNN网络通过可学习的参数处理输入的上一层特征图信息、上一层输出的粗糙光流（金字塔顶层除外）和相关性匹配层输出的特征图相关性信息，输出当前层的细化后的光流，细化后的光流和特征图经过反卷积层放大至自身两倍，作为下一层的输入；
- **Context Network：**Context Network由CNN网络组成，对最终输出的光流进行优化，其中CNN网络使用空洞卷积，在不引入额外计算量的前提下提高感受野，优化光流质量。
参考资料：
2.1 模型改动点
相比于公版模型，地平线所做的改动主要为：
- 出于模型优化考虑，在Backbone以及Head的部分卷积模块中，激活函数替换LeakyReLU为ReLU，并使用Conv+BN+ReLU的结构，其中BN模块对于量化更加友好；
- 在Warp操作中，去掉了地平线不支持的公版mask相关操作。
2.2 源码说明
2.2.1 Config文件
configs/opticalflow_pred/pwcnet/pwcnet_pwcnetneck_flyingchairs.py为该模型的配置文件，定义了模型结构和数据集加载，以及整套训练流程，所需参数的说明在算子定义中会给出。配置文件主要内容包括：
# 基础参数配置
task_name = "pwcnet_opticalflow"
batch_size_per_gpu = 64
device_ids = [0, 1, 2, 3, 4, 5, 6, 7] # 按实际GPU卡数修改
march = March.BAYES # J5平台
...
# 训练参数配置
num_train_samples = 22232
lr_base = 0.0001
max_steps = 1200000
max_epoch = int(max_steps * 16 / num_train_samples)
...
# 模型参数配置
out_channels = [16, 32, 64, 96, 128, 196] # Backbone每层特征图输出通道数
flow_pred_lvl = 2 # 光流最终输出层的序号
pyr_lvls = 6 # Backbone金字塔特征图层数
...
# 模型结构定义
model = dict(
type="PwcNet",
backbone=dict(
type="PwcNetNeck",
out_channels=out_channels,
pyr_lvls=pyr_lvls,
flow_pred_lvl=flow_pred_lvl,
act_type=nn.ReLU(), # 对应模型改动点1
...
),
head=dict(
type="PwcNetHead",
in_channels=out_channels,
md=4, # 计算Correlation时搜索框半径
pyr_lvls=pyr_lvls,
flow_pred_lvl=flow_pred_lvl,
act_type=nn.ReLU(), # 对应模型改动点1
...
),
loss=dict(type="LnNormLoss", norm_order=2, power=1, reduction="mean"), # L2范数Loss
...
)
deploy_model = dict(
...
loss=None,
)
deploy_inputs = dict(img=torch.randn((1, 6, 384, 512))) # 模型输入尺寸1x6x384x512
# 数据加载
data_loader = dict(
type=torch.utils.data.DataLoader,
dataset=dict(
type="FlyingChairs",
data_path="./tmp_data/FlyingChairs/train_lmdb/", # 将数据集处理成lmdb格式且根据实际目录修改
...
),
batch_size=batch_size_per_gpu,
...
)
val_data_loader = dict(...)
calibration_data_loader = copy.deepcopy(data_loader)
qat_data_loader = copy.deepcopy(data_loader)
align_bpu_data_loader = dict(...) # 模型仿真板端运行数据集加载
# Callback
stat_callback = dict(...)
val_callback = dict(...)
ckpt_callback = dict(...)
trace_callback = dict(...)
# 不同stage的训练策略配置
float_trainer = dict(...)
calibration_trainer = dict(...) # calibration_step = 10
qat_trainer = dict(...)
int_infer_trainer = dict(...) # num_epochs = 0，仅为保存int_infer pth & pt模型
# 不同stage的验证
float_predictor = dict(...)
calibration_predictor = dict(...)
qat_predictor = dict(...)
int_infer_predictor = dict(...)
align_bpu_predictor = dict(...) # 模型仿真板端运行验证
# 编译配置
compile_dir = os.path.join(ckpt_dir, "compile")
compile_cfg = dict(
march=march,
...
)
# 模型推理结果可视化—输入前处理
def process_inputs(...):
...
# 模型推理结果可视化—输出后处理
def process_outputs(...):
...
# 模型推理结果可视化配置
infer_cfg = dict(...)
# ONNX导出配置
onnx_cfg = dict(
model=deploy_model,
...
)
注：如果需要复现精度，config中的训练策略最好不要修改，否则可能会有意外的训练情况出现。
2.2.2 模型结构
地平线光流估计PWCNet参考算法的模型结构如下所示：
2.2.2.1 Backbone
-
Backbone由3x3的二维卷积所组成的多层CNN网络来提取输入图像的多尺度金字塔特征。输入尺寸1x6x384x512为两帧1x3x384x512的图像堆叠而成，并在后面的处理中拆分，共享一个网络，分别提取特征，达到了不引入额外参数量的目的。
class PwcNetNeck(nn.Module):
...
self.pyr_lvls = pyr_lvls # 特征图金字塔层数
self.flow_pred_lvl = flow_pred_lvl # 光流最终输出层的序号
self.out_channels = out_channels # 每层输出的特征图通道数
# 构建特征提取网络
self.mod = nn.ModuleList()
self.mod.append(self._make_stage(3, self.out_channels[0], act_type))
for idx in range(self.pyr_lvls - 1):
self.mod.append(
self._make_stage(
self.out_channels[idx],
self.out_channels[idx + 1],
act_type,
)
)
...
# 构建网络每一层，使用3x3的二维卷积，替换LeakyReLU为ReLU
def _make_stage(self, in_channels, out_channels, act_type):
layers = []
layers.append(
ConvModule2d(...),
)
for _ in range(2):
layers.append(
ConvModule2d(...),
)
return nn.Sequential(*layers)
# 前向推理
def forward(self, x):
x = self.quant(x)
output1 = []
output2 = []
x1 = x[:, :3, ...]
x2 = x[:, 3:, ...] # 对输入进行拆分
for module in self.mod:
x1 = module(x1)
x2 = module(x2) # 拆分后的输入共享一个网络，分别提取特征
output1.append(x1)
output2.append(x2)
return [output1, output2]
其中_make_stage在构建每层网络时的具体模块使用以及地平线参考算法对公版模型所做的改动如下：
对应代码路径：
/usr/local/lib/python3.8/dist-packages/hat/models/task_modules/pwcnet/neck.py
2.2.2.2 Head
Head部分对Backbone提取的多尺度金字塔特征进行处理，通过相关性匹配和Warp映射操作实现输出光流的Coarse-to-fine，并最终输出Refined光流。根据模型结构图可知Head部分主要组成模块为：Correlation（相关性匹配）、Warp（光流映射）和PredFlow（光流估计）。
class PwcNetHead(nn.Module):
...
self.in_channels = in_channels # Backbone输出的多尺度特征图通道数
self.pyr_lvls = pyr_lvls # 特征图金字塔层数
self.flow_pred_lvl = flow_pred_lvl # 光流最终输出层的序号
self.act_type = act_type # 使用ReLU作为激活函数
predict_flow_channels = [128, 128, 96, 64, 32] # PredFlow模块输出的特征图通道数
corr_out_channel = (md * 2 + 1) ** 2 # Correlation模块输出的通道数
# PredFlow
self.predict_flow = nn.ModuleList()
...
for idx in range(pyr_lvls, flow_pred_lvl - 1, -1):
...
self.predict_flow.append(
PredFlow(...)
)
...
# Correlation
self.corr = hpp.nn.Correlation(max_displacement=md, pad_size=md)
...
# Warp
def warp(...):
...
# 前向推理
def forward(...):
feat0 = features[0] # 第一张图I1的多尺度特征
feat1 = features[1] # 第二张图I2的多尺度特征
# 光流Coarse-to-fine
flow_list = []
for idx in range(self.pyr_lvls, self.flow_pred_lvl - 1, -1):
if idx == self.pyr_lvls: # 顶层特征图处理
corr = self.corr(feat0[idx - 1], feat1[idx - 1]) # 相关性计算
corr = self.act_type(corr)
flow, up_feat, up_flow = self.predict_flow[
self.pyr_lvls - idx
](corr) # 光流估计
flow_list.append(flow)
else: # 非顶层特征图处理
feat1_warp = self.warp(feat1[idx - 1], up_flow, idx) # 光流映射
corr = self.corr(feat0[idx - 1], feat1_warp)
corr = self.act_type(corr)
corr = self.concat[self.pyr_lvls - idx - 1].cat(
(corr, feat0[idx - 1], up_flow, up_feat), 1
) # 在C维度整合上层粗糙光流信息、上层与本层特征图信息和相关性信息
if idx == self.flow_pred_lvl: # 光流最终输出层
flow = self.predict_flow[self.pyr_lvls - idx](corr)
flow_list.append(flow)
else: # 非光流最终输出层，继续向下一层输出本层特征信息
flow, up_feat, up_flow = self.predict_flow[
self.pyr_lvls - idx
](corr)
flow_list.append(flow)
return self._dequant(flow_list) # 对输出进行反量化操作
对应代码路径：/usr/local/lib/python3.8/dist-packages/hat/models/task_modules/pwcnet/head.py
2.2.2.2.1 Correlation
Correlation部分进行特征图的相关性匹配，也就是公版模型提到的Cost Volume，其计算过程和卷积类似，都是对应位置特征值的点积和，不同之处在于特征图Correlation的Kernel为另一张特征图或另一张特征图的patch，且不同于卷积操作需要先对Kernel进行翻转。公式如下：
c(x1,x2)=∑o∈[−k,k]×[−k,k]<f1(x1+o),f2(x2+o)>c(x_{1}, x_{2}) = \sum_{o \in [-k,k] \times [-k,k]} <f_{1}(x_{1} + o), f_{2}(x_{2} + o)>c(x1,x2)=o∈[−k,k]×[−k,k]∑<f1(x1+o),f2(x2+o)>
其中k=2D+1k=2D+1k=2D+1为Correlation Kernel的边长，D=4D=4D=4为本模型所取的Kernel半径，和公版一致。Correlation部分无可训练的权重参数，减少了参数量从而达到模型轻量化的目标。
class Correlation(nn.Module):
def __init__(
self,
kernel_size=1,
max_displacement=4, # 确定data2中搜索范围的最大偏移量
stride1=1, # data1的全局滑动步长
stride2=1, # data2搜索框内的滑动步长
pad_size=4,
is_multiply=True, # 使用乘积和计算Correlation
):
...
self.kernel_radius = (self.kernel_size - 1) / 2 # 乘积和对象为单个特征值，尺寸为1x1
self.border_size = self.kernel_radius + self.max_displacement # 边界大小（最大偏移量+kernel半径） = 4
self.neighborhood_grid_radius = self.max_displacement // self.stride2 # 搜索框半径 = 4
self.neighborhood_grid_width = self.neighborhood_grid_radius * 2 + 1 # 搜索框边长D = 9
self.top_channels = int(
self.neighborhood_grid_width * self.neighborhood_grid_width
) # 输出通道数D**2 = 81，data2[0:8, 0:8]
...
def _forward(...):
...
top_height = math.ceil(
float(height + 2 * pad - self.border_size * 2) / float(strides[0])
)
top_width = math.ceil(
float(width + 2 * pad - self.border_size * 2) / float(strides[0])
) # 确定Correlation Map输出尺寸，top_height, top_width = H, W
...
inter_height = (top_height - 1) * strides[0] + kernel
inter_width = (top_width - 1) * strides[0] + kernel # 确定inter_out输出尺寸，inter_height, inter_width = H, W
tmp1 = nn.functional.pad(data1, (pad,) * 4).to(data1.device)
tmp2 = torch.zeros_like(tmp1, device=data1.device)
tmp2[
:,
:,
pad : pad + data2.shape[2],
pad : pad + data2.shape[3],
] = data2 # tmp1为padding后data1，tmp2为padding后data2
out = torch.zeros((num, self.top_channels, top_height, top_width)).to(
data1.device
) # 初始化Correlation Map输出：[N,81,H,W]
inter_out = torch.zeros(
(num, self.top_channels, inter_height, inter_width)
).to(data1.device) # 初始化inter_out输出：[N,81,H,W]
self.sumelems = kernel * kernel * channels # kernel内特征值数量，sumelems = 1 * 1 * C
# 计算inter_out输出
for c in range(self.top_channels):
x1 = max_displacement
y1 = max_displacement
s2o = int(...)
s2p = int(...)
x2 = max_displacement + s2o
y2 = max_displacement + s2p # (y1, x1)和(y2, x2)分别为tmp1和tmp2目标区域top_left点坐标(h, w)
if is_multiply:
inter_out[:, c, :, :] = torch.sum(
tmp1[
:, :, y1 : y1 + inter_height, x1 : x1 + inter_width
].mul(
tmp2[
:, :, y2 : y2 + inter_height, x2 : x2 + inter_width
]
),
(1),
) # 合并C维度的Correlation乘积和
...
...
# 计算Correlation Map输出
for i in range(top_height):
for j in range(top_width):
y1 = i * strides[0]
x1 = j * strides[0]
out[:, :, i, j] = torch.sum(
inter_out[
:,
:,
y1 : y1 + kernel,
x1 : x1 + kernel,
],
(2, 3),
) # H和W维度上合并Correlation计算结果
out /= self.sumelems # 对kernel内特征值数取平均
return out
def forward(self, data1, data2):
"""Forward for Horizon Correlation:
data1: [N,C,H,W]
data2: [N,C,H,W]
Return:
output: [N,81,H,W]
"""
return self._forward(...)
对应代码路径：/usr/local/lib/python3.8/dist-packages/horizon_plugin_pytorch/nn/correlation.py
2.2.2.2.2 Warp
Warp部分通过上一层输出的光流估计（顶层除外），对第二帧图像的对应特征图进行光流映射操作，其设计的基础思想是引用了Lucas-Kanade算法（稀疏光流估计代表算法，详细算法介绍可参考光流Lucas-Kanade算法）中的一个重要假设：亮度恒定。根据亮度恒定假设得出的约束方程如下：
I(x,y,t)=I(x+δx,y+δy,t+δt)I(x,y,t)=I(x+\delta{x},y+\delta{y},t+\delta{t})I(x,y,t)=I(x+δx,y+δy,t+δt)
I(x,y,t)I(x,y,t)I(x,y,t)表示ttt时刻(x,y)(x,y)(x,y)位置像素点的强度，或是特征点的特征值。因此如果光流(δx,δy)(\delta{x},\delta{y})(δx,δy)估计越准确，第二帧图像映射后的特征图和第一帧图像对应层级的特征图相关性越高，对应Correlation计算结果越大。-
对于顶层特征图的处理，公版是手动构造全“0”元素的光流图对顶层特征图进行映射，而地平线参考算法是直接跳过顶层特征图的Warp操作，两者结果相同；此外相比公版实现，地平线去掉了不支持的公版mask相关操作，公版代码参考：PWCNet公版Warp源码。Warp操作通过grid_sample算子实现，并且无可训练参数，减小了模型尺寸。
def warp(self, x: torch.Tensor, up_flow: torch.Tensor, idx: int) -> torch.Tensor:
"""基于上一层输出的光流估计，将I2图像特征图映射回I1图像：
x: [N,C,H,W] — I2图像特征图
up_flow: [N,2,H,W] — 上一层输出的光流估计
idx — 当前层序号
Return:
output: [N,2,H,W] — 映射后的特征图
"""
flo = up_flow.mul(self.warp_scale[self.pyr_lvls - idx])
vgrid = flo.permute(0, 2, 3, 1)
output = self.grid_sample(x, vgrid)
return output
对应代码路径：/usr/local/lib/python3.8/dist-packages/hat/models/task_modules/pwcnet/head.py
2.2.2.2.3 PredFlow
PredFlow部分输入为在通道维度上整合的粗糙光流信息、特征图信息和相关性信息，使用一个CNN网络处理并输出估计的光流，以及通过反卷积层对本层估计的光流和特征图进行处理，为下一层的光流估计提供信息，实现光流的Coarse-to-fine。在CNN网络中引入了Dense Connection和Residual Connection，进行特征复用以及获得更好的梯度流信息，提高模型的精度和泛化能力。此外引入Context Network对光流进行进一步优化。
class PredFlow(nn.Module):
...
self.in_channel = in_channel # [81, 81+2+2+128, 81+2+2+96, 81+2+2+64, 81+2+2+32]
self.out_channles = out_channles # [128, 128, 96, 64, 32]
self.act_type = act_type # 使用ReLU为激活函数
...
self.mod = self._make_mod(self.in_channel, self.out_channles) # CNN网络
...
self.pred_flow = nn.Sequential(ConvModule2d(...),) # 对输出光流进行Refine，act_layer = None
...
self.deconv_flow = ConvTransposeModule2d(...) # 光流反卷积层，C_IN, C_OUT = 2, 2
self.deconv_feat = ConvTransposeModule2d(...) # 特征图反卷积层，C_IN, C_OUT = out_mod_channels, 2
# 构建CNN网络
def _make_mod(self, in_channel, out_channles):
mod = nn.ModuleList()
for idx in range(len(out_channles)):
mod.append(
nn.Sequential(
ConvModule2d(
...
)
)
)
...
# Context Network
def _make_res_cxt(self, in_channel, out_channels, dilations):
...
...
# 前向推理
def forward(self, x: torch.Tensor) -> List[torch.Tensor]:
for idx, module in enumerate(self.mod):
x1 = module(x)
if self.use_dense: # 引入Dense Connection，提高模型精度和泛化能力
x = self.concat[idx].cat((x1, x), 1)
else:
x = x1
flow = self.pred_flow(x) # 光流Refine
if self.use_res: # 引入Residual Connection，并使用Context Network优化光流
flow = self.short_add.add(flow, self.res_cxt(x))
if self.is_pred_lvl: # 光流最终输出层
return flow
up_feat = self.deconv_feat(x) # 对特征图进行反卷积，尺寸x2，输出到下一层处理
up_flow = self.deconv_flow(flow) # 对光流进行反卷积，尺寸x2，输出到下一层处理
return [flow, up_feat, up_flow]
对应代码路径：/usr/local/lib/python3.8/dist-packages/hat/models/task_modules/pwcnet/head.py
2.2.3 Loss
Training Loss用来衡量每一层估计的光流（共5层）和光流真值之间的距离，也就是公版论文提到的多尺度光流Training Loss，公式如下：
-
Θ\ThetaΘ表示可学习参数的集合，xxx表示像素点或特征点索引，公式中∣⋅∣2|\cdot|_{2}∣⋅∣2表示计算L2L2L2范数，wΘl(x)w^{l}_{\Theta}(x)wΘl(x)表示lll层的估计光流，wGTl(x)w^{l}_{GT}(x)wGTl(x)表示lll层的光流真值，αl\alpha_{l}αl为lll层的Loss加权值，γ\gammaγ为正则化参数。
class LnNormLoss(nn.Module):
"""LnNorm loss:
norm_order: 2 — L2 Norm
epsilon: 0 — A small constant for finetune
power: 1 — A power num of norm + epsilon of loss
reduction: mean — Reduction mode
loss_weight: [0.005, 0.01, 0.02, 0.08, 0.32] — Used to weight the output
"""
...
def forward(
self,
pred: torch.Tensor, # 估计的光流，[N, 2, H, W]
target: torch.Tensor, # 光流真值，通过对数据集真值采样得到，[N, 2, H, W]
weight: Optional[torch.Tensor] = None, # loss_weight
avg_factor: Optional[Union[float, torch.Tensor]] = None, # Normalized factor
) -> torch.Tensor:
diff = pred - target.detach()
loss = torch.norm(diff, p=self.norm_order, dim=1)
loss = loss.sum((1, 2))
loss = loss.mean()
loss = loss + self.epsilon
loss = torch.pow(loss, self.power)
loss = weight_reduce_loss(loss, weight, self.reduction, avg_factor)
if self.loss_weight is not None:
loss *= self.loss_weight
return loss
对应代码路径：/usr/local/lib/python3.8/dist-packages/hat/models/losses/lnnorm_loss.py
2.2.4 EPE Metric
EPE（Endpoint Error）是评估光流估计算法精度的指标，计算方法为估计的光流向量终点二维坐标和真值光流向量终点二维坐标的欧式距离。
class EndPointError(EvalMetric):
"""Metric for OpticalFlow task, Endpoint Error (EPE)
"""
...
def update(self, labels, preds, masks=None):
if self.use_mask: # self.use_mask = False，对应改动点2
...
else:
diff = preds - labels
bs = preds.shape[0]
epe = torch.norm(diff, p=2, dim=1).mean((1, 2)).sum().item()
self.sum_metric += epe
self.num_inst += bs
对应代码路径：/usr/local/lib/python3.8/dist-packages/hat/metrics/metric_optical_flow.py
2.2.5 后处理
光流估计后处理分为两部分：第一步将输出的光流图通过双线性插值操作恢复到原图分辨率；第二步进行光流图的可视化，可视化具体做法为用色相（Hue）去表示光流向量的方向，用色度（Saturation）去表示光流向量的长度，简而言之，就是用不同的颜色去表示像素点不同的运动方向，像素点运动偏移越大，该色彩的亮度越高。
def process_outputs(model_outs, viz_func, vis_inputs):
preds = model_outs
preds = F.interpolate(preds.float(), scale_factor=4, mode="bilinear") * 4.0 # 第一步恢复分辨率
preds = preds.permute((0, 2, 3, 1))
preds = viz_func(vis_inputs, preds) # 第二步可视化
return None
对应代码路径：configs/opticalflow_pred/pwcnet/pwcnet_pwcnetneck_flyingchairs.py-
可视化代码路径：/usr/local/lib/python3.8/dist-packages/hat/visualize/opticalflow.py
3 浮点模型训练
3.1 Before Start
3.1.1 发布物及环境部署
step1: 获取发布物-
J5：下载OE包horizon_j5_open_explorer_v$version$-py38_$releasedate$.tar.gz，获取方式见地平线开发者社区【OpenExplorer J5算法工具链 版本发布】。-
step2: 解压发布包-
获取发布包后，首先对其进行解压：
tar -zxvf horizon_j5_open_explorer_v$version$-py38_$releasedate$.tar.gz
解压后文件结构如下：
|-- bsp
|-- ddk
| |-- package
| `-- samples
| |-- ai_benchmark
| |-- ai_forward_view_sample
| |-- ai_toolchain
| | |-- ...
| | |-- horizon_model_train_sample # 参考算法目录
| | `-- model_zoo
| |-- model_zoo
| `-- vdsp_rpc_sample
|-- README-CN
|-- README-EN
|-- resolve_all.sh
`-- run_docker.sh
其中horizon_model_train_sample为参考算法模块，包含以下模块：
|-- horizon_model_train_sample #参考算法示例
| |-- plugin_basic #qat 基础示例
| `-- scripts #模型配置文件
step3: 拉取GPU docker环境
docker pull openexplorer/ai_toolchain_ubuntu_20_j5_gpu:"$version"
#启动容器，具体参数可根据实际需求配置
#-v 用于将本地的路径挂载到 docker 路径下
nvidia-docker run -it --shm-size="15g" -v `pwd`:/WORKSPACE openexplorer/ai_toolchain_ubuntu_20_j5_gpu:"$version"
3.1.2 数据集准备
3.1.2.1 数据集下载
下载官方的FlyingChairs数据集作为网络的训练集和验证集，下载链接为FlyingChairs.zip，同时需要下载相应的标签数据FlyingChairs_train_val.txt，解压缩之后数据目录结构如下所示：
# horizon_model_train_sample/scripts/tmp_data
tmp_data
|-- FlyingChairs
|-- FlyingChairs_release
|-- data
|-- README.txt
|-- FlyingChairs_train_val.txt
|-- FlyingChairs.zip
3.1.2.2 数据集打包
为了提升训练速度，我们需要对原始格式的数据集进行打包，转换成lmdb格式的数据集。进入horizon_model_train_sample/scripts目录，使用以下命令将训练数据集和验证数据集打包，格式为lmdb：
#打包训练数据集
python3 tools/datasets/flyingchairs_packer.py --src-data-dir ./tmp_data/FlyingChairs --split-name train --pack-type lmdb --num-workers 10 --target-data-dir ./tmp_data/FlyingChairs
#打包验证数据集
python3 tools/datasets/flyingchairs_packer.py --src-data-dir ./tmp_data/FlyingChairs --split-name val --pack-type lmdb --num-workers 10 --target-data-dir ./tmp_data/FlyingChairs
--src-data-dir为解压后的FlyingChairs数据集目录；-
--target-data-dir为打包后数据集的存储目录。
数据集打包命令执行完毕后会在target-data-dir下生成train_lmdb和val_lmdb，train_lmdb和val_lmdb就是打包之后的训练数据集和验证数据集，也就是configs/opticalflow_pred/pwcnet/pwcnet_pwcnetneck_flyingchairs.py文件中的data_path，打包后的目录结构如下：
# horizon_model_train_sample/scripts/tmp_data
tmp_data
|-- FlyingChairs
|-- FlyingChairs_release
|-- data
|-- README.txt
|-- FlyingChairs_train_val.txt
|-- FlyingChairs.zip
|-- train_lmdb
|-- val_lmdb
3.1.3 Config配置
在进行模型训练和验证之前，需要对configs/opticalflow_pred/pwcnet/pwcnet_pwcnetneck_flyingchairs.py文件中的部分参数进行配置，一般情况下，我们需要配置以下参数：
- device_ids、batch_size_per_gpu：根据实际硬件配置进行gpu id和每个gpu的batchsize的配置；
- ckpt_dir：浮点、calib、量化训练的权重路径配置，权重下载链接在对应
configs/opticalflow_pred/pwcnet/README.md文件中； - data_path：数据集的存放路径，按照实际存放路径配置。
3.2 浮点模型训练
环境部署、数据集准备和config文件中的参数配置完成后，使用以下命令训练浮点模型：
python3 tools/train.py --config configs/opticalflow_pred/pwcnet/pwcnet_pwcnetneck_flyingchairs.py --stage float
浮点模型训练后模型ckpt的保存路径为config配置的
ckpt_callback中save_dir的值，默认为ckpt_dir。
3.3 浮点模型精度验证
浮点模型训练完成以后，可以使用以下命令验证已经训练好的浮点模型精度：
python3 tools/predict.py --config configs/opticalflow_pred/pwcnet/pwcnet_pwcnetneck_flyingchairs.py --stage float
验证完成后，会在终端打印浮点模型在验证集上的精度，如下所示：
...
INFO [metric_updater.py:360] Node[0] Epoch[0] Validation pwcnet_opticalflow: EPE[1.4114]
INFO [logger.py:176] Node[0] ==================================================END PREDICT==================================================
INFO [logger.py:176] Node[0] ==================================================END FLOAT PREDICT==================================================
4 模型量化和编译
完成浮点训练后，还需要进行量化训练和编译，才能将定点模型部署到板端。地平线对该模型的量化采用horizon_plugin框架，经过Calibration+QAT量化训练后，使用compile工具将量化模型编译成可以上板运行的“hbm”文件。
4.1 Calibration
为加速QAT训练收敛和获得最优量化精度，建议在QAT训练之前做Calibration，其过程为通过指定的样本初始化量化参数，为QAT量化训练提供一个更好的初始化参数。Calibration和浮点训练的方式相似，将checkpoint_path指定为训好的浮点权重路径。通过运行下面的命令就可以开启模型的Calibration：
python3 tools/train.py --config configs/opticalflow_pred/pwcnet/pwcnet_pwcnetneck_flyingchairs.py --
stage calibration
Calibration完成后会在终端打印保存checkpoint的相关信息。
4.2 Calibration模型精度验证
Calibration完成以后，可以使用以下命令验证Calibration模型的精度：
python3 tools/predict.py --config configs/opticalflow_pred/pwcnet/pwcnet_pwcnetneck_flyingchairs.py --stage calibration
验证完成后，会在终端输出Calibration模型在验证集上的精度，格式见3.3。
4.3 量化模型训练
Calibration完成后，就可以加载calib权重开启模型的量化训练。量化训练其实是在浮点训练基础上的fine-tune，具体配置信息在config的qat_trainer中定义。量化训练的时候，初始学习率设置为浮点训练的十分之一，训练的epoch次数也大大减少。和浮点训练的方式相似，将checkpoint_path指定为训好的calib权重路径。通过运行下面的命令就可以开启模型的量化训练：
python3 tools/train.py --config configs/opticalflow_pred/pwcnet/pwcnet_pwcnetneck_flyingchairs.py --stage qat
量化训练后模型ckpt的保存路径为config配置的ckpt_callback中的save_dir的值，默认为ckpt_dir。
4.4 量化模型精度验证
量化训练完成后，通过运行以下命令进行量化模型的精度验证：
#qat模型精度验证
python3 tools/predict.py --config configs/opticalflow_pred/pwcnet/pwcnet_pwcnetneck_flyingchairs.py --stage qat
#quantized模型精度验证
python3 tools/predict.py --config configs/opticalflow_pred/pwcnet/pwcnet_pwcnetneck_flyingchairs.py --stage int_infer
qat模型的精度验证对象为插入伪量化节点后的模型（float32）；quantized模型的精度验证对象为定点模型（int8），验证的精度是最终的int8模型的真正精度，这两个精度应该是十分接近的。
验证完成后，会在终端输出qat模型和定点模型在验证集上的精度，格式见3.3。
4.5 仿真精度验证
除了上述模型精度验证之外，我们还提供仿真“hbm”模型上板的精度验证方法，可以通过下面的命令完成：
python3 tools/align_bpu_validation.py --config configs/opticalflow_pred/pwcnet/pwcnet_pwcnetneck_flyingchairs.py
验证完成后，会在终端输出仿真上板的精度，格式见3.3。
4.6 量化模型编译
在量化训练完成之后，可以使用“compile_perf.py”脚本将量化模型编译成可以板端运行的“hbm”模型，同时该工具也能预估模型在BPU上的运行性能，“compile_perf”脚本使用方式如下：
python3 tools/compile_perf.py --config configs/opticalflow_pred/pwcnet/pwcnet_pwcnetneck_flyingchairs.py --opt 3
opt为优化等级，取值范围为0~3，数字越大优化等级越高，编译时间也会越长；-
可以指定--out_dir为编译后产出物的存放路径，默认在ckpt_dir的compile文件夹下。
运行后，ckpt_dir的compile目录下会产出以下文件：
|-- compile
|-- .html #模型在bpu上的静态性能数据
|-- .json
|-- model.hbm #用于板端部署的模型
|-- model.hbir #编译过程的中间文件
`-- model.pt #模型的pt文件
5 其他工具
5.1 导出ONNX模型
我们提供将定点模型导出为ONNX格式的脚本，使用方法如下：
python3 tools/export_onnx.py --config configs/opticalflow_pred/pwcnet/pwcnet_pwcnetneck_flyingchairs.py
运行结束会在
ckpt_dir文件夹下生成产物qat.onnx。
5.2 结果可视化
如果你希望可以看到训练出来的模型推理结果的可视化效果，我们在tools文件夹下提供了模型推理及可视化的脚本，你需要先修改configs/opticalflow_pred/pwcnet/pwcnet_pwcnetneck_flyingchairs.py中的infer_cfg：
infer_cfg = dict(
model=model,
infer_inputs=dict(
img1="./tmp_data/FlyingChairs/FlyingChairs_release/data/00411_img1.ppm", #修改为自己的路径
img2="./tmp_data/FlyingChairs/FlyingChairs_release/data/00411_img2.ppm", #修改为自己的路径
),
...
)
然后运行以下命令即可：
python3 tools/infer.py --config configs/opticalflow_pred/pwcnet/pwcnet_pwcnetneck_flyingchairs.py --save-path ./
--save-path:可视化结果保存路径
可视化示例：
6 板端部署
6.1 上板性能实测
使用hrt_model_exec perf工具将生成的“hbm”文件上板做BPU性能的FPS实测，hrt_model_exec perf参数如下：
hrt_model_exec perf --model_file {model}.hbm \
--thread_num 8 \
--frame_count 1000 \
--core_id 0 \
--profile_path './'
6.2 AI Benchmark示例
OE开发包中提供了光流估计参考算法的AI Benchmark示例，位于ddk/samples/ai_benchmark/j5/qat/script/opticalflow/pwcnet_opticalflow。示例的具体使用方法可以参考开发者社区J5算法工具链产品手册—AI Benchmark评测示例，在板端执行以下命令进行模型性能评测：
#帧率数据
sh fps.sh
#单帧延迟数据
sh latency.sh
如果要进行模型精度评测，请参考开发者社区J5算法工具链产品手册—AI Benchmark示例精度评测进行数据的准备和模型的推理。