1 简介
本文提出了YolactEdge实时实例分割方法，可以在小型边缘设备上以实时速度运行。具体来说，在550x550分辨率的图像上，带有ResNet-101主干的YolactEdge在Jetson AGX Xavier上的运行速度高达30.8FPS（在RTX 2080Ti上的运行速度为172.7FPS）。为了实现这一目标，我们对基于图像的最新实时方法YOLACT进行了两项改进：
1）优化TensorRT，同时谨慎权衡速度和准确性；
2）利用视频中时间冗余的新型特征扭曲模块。
在YouTube VIS和MS COCO数据集上进行的实验表明，与现有的实时方法相比，YolactEdge的速度提高了3-5倍，同时具有极好的mask和box检测精度。
2 相关工作
YOLACT是第一个实时实例分割方法，以达到具有挑战性的MS COCO数据集的竞争精度。最近，CenterMask、BlendMask和SOLOv2通过利用更精确的目标探测器(如FCOS)，在一定程度上提高了精度。所有现有的实时实例分割方法都是基于图像的，并且需要像Titan Xp/RTX 2080Ti这样庞大的gpu来实现实时速度。相反，本文提出了第一个基于视频的实时实例分割方法，可以在小边缘设备上运行，如Jetson AGX Xavier。
2.1 视频中的特征传播
很多方法将视频中的特征传播用于提高视频分类和视频目标检测的速度和准确性。这些方法使用现成的光流网络来估计像素级的物体运动，并在帧与帧之间扭曲特征图。然而，即使是最轻量级的流网络，也需要不可忽略的内存和计算，这是阻碍边缘设备实时速度的障碍。相比之下，本文的模型在特征级别(与输入像素级别相反)直接估计物体运动并执行特征扭曲，从而实现实时速度。
2.2 提高模型效率
设计轻量级的高性能骨干和特征金字塔已经成为提高深度网络效率的主要动力之一。MobileNetv2引入了深度卷积和反向残差，为移动设备设计了一个轻量级架构。MobileNetv3、nas-fpn和EfficientNet使用神经结构搜索自动找到高效的结构。其他的则利用知识蒸馏、模型压缩或二进位网络。CVPR低功耗计算机视觉挑战参与者使用TensorRT量化和加速目标检测器，如NVIDIA Jetson TX2上的Fast-RCNN。与这些方法相比，YolactEdge保留了庞大的Backbone，并利用了视频中的时间冗余，同时使用TensorRT优化实现了快速、准确的实例分割。
3 本文方法
3.1 TensorRT优化
表1显示了对YOLACT的分析，这是YolactEdge的基线模型。
简单地说，YOLACT可分为4个部分:
- 1)特征主干;
- 2)特征金字塔网络(FPN);
- 3)ProtoNet;
- 4)预测头;
网络架构如图1(右)所示。上表I的第2行代表YOLACT，所有组件都在FP32中(即，没有进行TensorRT优化)，结果在Jetson AGX Xavier上只有6.6FPS，带有ResNet-101主干网。INT8或FP16转换在不同的模型组件导致各种改进的速度和精度的变化。值得注意的是，将预测头转换为INT8(最后四行)总是会导致大量实例分割精度的损失。
假设这是因为最终的box和mask预测需要在最终的表示中不丢失，同时需要超过个Boxes才能被编码。将除预测头和FPN(灰色突出显示的行)以外的所有组件转换为INT8可以获得最高的FPS，且map退化很小。因此，这是在实验中为模型使用的最终配置，但是可以根据需要轻松地选择不同的配置。
为了量化模型元件到INT8精度，校准步骤是必要的。为此，TensorRT收集每一层激活的直方图生成几个具有不同阈值的量化分布，并使用KL散度将每个量化分布与参考分布进行比较。这个步骤确保模型在转换为INT8精度时损失尽可能小的性能。
上表显示校准数据集大小的影响。文章实验发现使用50或100张图像进行校准在精度和速度方面都足够了。
3.2 利用视频中的时间冗余
TensorRT的优化使速度有了约4x的提升，在处理静态图像时，应该使用YolactEdge的这个版本。然而，在处理视频时可以利用时间冗余来做。
这里给出一个序列帧的视频，目标是在每一帧中以快速和准确的方式预测每个对象实例的mask。对于视频实例分割网络，大体上遵循YOLACT速度和准确性权衡的设计理念。具体地说，在每一帧上执行2项并行任务:
- 1)生成一组原型mask；
- 2)预测每个实例的掩模系数。然后，将原型与掩模系数线性组合，得到最终的掩模。
为了清晰地表示，将分解为和，其中表示特征骨干阶段，表示剩余部分(即用于生成原型掩模的类、box、mask系数和ProtoNet的预测头)，它们接受的输出并进行实例分割预测。
本文有选择地将视频中的帧分为2类:关键帧和非关键帧；模型在这2组框架上只是在Backbone阶段发生了变化。
对于关键帧,模型计算所有Backbone和金字塔特征(在Fig 1中的C1−C5和P3−P7)。而对于非关键帧只计算一个子集，把其余的利用时间冗余机制进行计算。
通过这种方式可以在生成准确预测的同时保持快速运行时之间取得平衡。
3.3 部分功能转换
从邻近关键帧转换(即扭曲)特征被证明是一种有效的策略，以减少主干计算，以产生快速的视频Box目标检测器。具体地说，是使用离网光流网络转换所有主干特征。然而，由于光流估计中不可避免的误差，我们发现它不能提供像实例分割这样的像素级任务所需的足够精确的特征。在本研究中提出执行部分特征转换，以改善转换后的特征的品质，同时仍维持快速的运行时间。
本文方法计算的非关键帧只有通过高分辨率级别(例如,跳过, 因此, 计算),且仅将低分辨率转换/特性从之前的关键帧近似, (表示/)在当前非关键帧——如图1所示(右)。它通过与YOLACT相同的方式对进行向下采样来计算/。通过计算特征和变换后的特征，生成为，其中表示上采样。
最后，使用特征来生成像素精确的原型。通过这种方式，可以保留用于生成mask原型的高分辨率细节，因为高分辨率的特征是计算而不是转换的，因此不会受到流量估计中的错误的影响。
重要的是，尽管为每一帧(即关键帧和非关键帧)计算骨干特征，但避免了计算backbone特征中最复杂的部分，因为金字塔网络的不同阶段的计算代价是高度不平衡的。
如表所示，ResNet-101花费66%以上的计算在，而一半以上的推理时间被Backbone计算占用。而通过只计算特征金字塔的底层，并转换其余的层可以大大加快速度以达到实时性能。
总而言之，部分特征转换设计产生了更高质量的特征图，同时也支持实时速度。
3.4 有效的运动估计
该部分描述如何有效地计算关键帧和非关键帧之间的流。
给定一个非关键帧及其前一关键帧，所设计模型首先将它们之间的物体运动编码为一个2-D流场。然后利用流场将特征从帧变换到与帧对齐得到warped特征。
为了进行快速的特征变换，需要有效地估计目标运动。现有的进行流导向特征变换的框架直接采用现成的像素级光流网络进行运动估计。例如，FlowNetS(Fig.2a)在3个阶段执行流量估计:
- 首先，以原始RGB帧作为输入并计算特征堆栈;
- 然后，它通过递归上采样和拼接特征映射来细化特征子集，生成既包含高级(大运动)信息又包含精细局部信息(小运动)的粗到细特征;
- 最后，它使用这些特征来预测最终的Flow Map。
在本算法中，为了节省计算成本没有使用现成的流网络来处理原始RGB帧，而是重新使用所设计骨干网络计算出的特征，它已经产生了一组语义丰富的特征。为此提出了FeatFlowNet(b)，它通常遵循FlowNetS架构，但在第1阶段，不是从原始RGB图像输入计算特征堆栈，而是重用来自ResNet骨干(C3)的特征，并使用较少的卷积层。正如实验中所证明的那样，Flow估计网络在同样有效的情况下要快得多。
class FlowNetMini(ScriptModuleWrapper): __constants__ = ['interpolation_mode', 'skip_flow'] def __init__(self, in_features): super().__init__() self.interpolation_mode = cfg.fpn.interpolation_mode self.conv1 = build_flow_convs(cfg.flow.encode_layers[0], in_features, cfg.flow.encode_channels, groups=cfg.flow.num_groups) self.conv2 = build_flow_convs(cfg.flow.encode_layers[1], cfg.flow.encode_channels, cfg.flow.encode_channels * 2, stride=2) self.conv3 = build_flow_convs(cfg.flow.encode_layers[2], cfg.flow.encode_channels * 2, cfg.flow.encode_channels * 4, stride=2) self.pred3 = FlowNetMiniPredLayer(cfg.flow.encode_channels * 4) self.pred2 = FlowNetMiniPredLayer(cfg.flow.encode_channels * 4 + 2) self.pred1 = FlowNetMiniPredLayer(cfg.flow.encode_channels * 2 + 2) self.upfeat3 = nn.Conv2d(cfg.flow.encode_channels * 4, cfg.flow.encode_channels * 2, kernel_size=3, stride=1, padding=1, bias=False) self.upflow3 = nn.Conv2d(2, 2, kernel_size=3, stride=1, padding=1, bias=False) self.upfeat2 = nn.Conv2d(cfg.flow.encode_channels * 4 + 2, cfg.flow.encode_channels, kernel_size=3, stride=1, padding=1, bias=False) self.upflow2 = nn.Conv2d(2, 2, kernel_size=3, stride=1, padding=1, bias=False) self.skip_flow = not self.training and cfg.flow.flow_layer != 'top' and '3' not in cfg.flow.warp_layers for m in chain(*[x.modules() for x in (self.conv1, self.conv2, self.conv3)]): if isinstance(m, nn.Conv2d) or isinstance(m, nn.ConvTranspose2d): nn.init.kaiming_normal_(m.weight.data, 0.1, mode='fan_in') if m.bias is not None: m.bias.data.zero_() @script_method_wrapper def forward(self, target_feat, source_feat): preds: List[Tuple[torch.Tensor, torch.Tensor, torch.Tensor]] = [] concat0 = shuffle_cat(target_feat, source_feat) out_conv1 = self.conv1(concat0) out_conv2 = self.conv2(out_conv1) out_conv3 = self.conv3(out_conv2) _, _, h2, w2 = out_conv2.size() flow3, scale3, bias3 = self.pred3(out_conv3) out_upfeat3 = F.interpolate(out_conv3, size=(h2, w2), mode=self.interpolation_mode, align_corners=False) out_upfeat3 = self.upfeat3(out_upfeat3) out_upflow3 = F.interpolate(flow3, size=(h2, w2), mode=self.interpolation_mode, align_corners=False) out_upflow3 = self.upflow3(out_upflow3) concat2 = torch.cat((out_conv2, out_upfeat3, out_upflow3), dim=1) flow2, scale2, bias2 = self.pred2(concat2) dummy_tensor = torch.tensor(0, dtype=out_conv2.dtype) if not self.skip_flow: _, _, h1, w1 = out_conv1.size() out_upfeat2 = F.interpolate(concat2, size=(h1, w1), mode=self.interpolation_mode, align_corners=False) out_upfeat2 = self.upfeat2(out_upfeat2) out_upflow2 = F.interpolate(flow2, size=(h1, w1), mode=self.interpolation_mode, align_corners=False) out_upflow2 = self.upflow2(out_upflow2) concat1 = torch.cat((out_conv1, out_upfeat2, out_upflow2), dim=1) flow1, scale1, bias1 = self.pred1(concat1) preds.append((flow1, scale1, bias1)) else: preds.append((dummy_tensor, dummy_tensor, dummy_tensor)) preds.append((flow2, scale2, bias2)) preds.append((flow3, scale3, bias3)) return preds
3.5 特征Warping
使用FeatFlowNet来估计先前关键帧和当前非关键帧之间的Flow图，然后通过反向warp将特征从转换为:
将中的每个像素投影到中为，其中。通过双线性插值计算像素值，其中为不同空间位置的双线性插值权值。
class FlowNetUnwrap(nn.Module): def forward(self, preds): outs: List[Tuple[torch.Tensor, torch.Tensor, torch.Tensor]] = [] flow1, scale1, bias1, flow2, scale2, bias2, flow3, scale3, bias3 = preds outs.append((flow1, scale1, bias1)) outs.append((flow2, scale2, bias2)) outs.append((flow3, scale3, bias3)) return outs class FlowNetMiniTRTWrapper(nn.Module): def __init__(self, flow_net): super().__init__() self.flow_net = flow_net if cfg.flow.use_shuffle_cat: self.cat = ShuffleCat() else: self.cat = Cat() self.unwrap = FlowNetUnwrap() def forward(self, a, b): concat = self.cat(a, b) dummy_tensor = torch.tensor(0, dtype=a.dtype) preds = [dummy_tensor, dummy_tensor, dummy_tensor] preds_ = self.flow_net(concat) preds.extend(preds_) outs = self.unwrap(preds) return outs
3.6 损失函数
对于实例分割任务，使用与YOLACT相同的损失来训练模型:
- 分类损失，
- Box回归损失，
- Mask损失，
- 以及辅助语义分割损失。
对于Flow估计网络的预训练使用端点误差(endpoint error, EPE)。
class OpticalFlowLoss(nn.Module): def __init__(self): super(OpticalFlowLoss, self).__init__() def forward(self, preds, gt): losses = {} loss_F = 0 for pred in preds: _, _, h, w = pred.size() gt_downsample = F.interpolate(gt, size=(h, w), mode='bilinear', align_corners=False) loss_F += torch.norm(pred - gt_downsample, dim=1).mean() losses['F'] = loss_F return losses
