学不动系列 | YolactEdge:边缘设备上的实时实例分割

YolactEdge：边缘设备上的实时实例分割（Xavier: 30 FPS, RTX 2080 Ti：170 FPS）带有ResNet-101的YolactEdge在Jetson AGX Xavier上的速度高达30.8 FPS，在RTX 2080 Ti上的速度为172.7 FPS，AP性能超强！速度是目前主流方法的3-5倍，代码刚刚开源！

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的这个版本。然而，在处理视频时可以利用时间冗余来做。

• 1)生成一组原型mask；
• 2)预测每个实例的掩模系数。然后，将原型与掩模系数线性组合，得到最终的掩模。

本文有选择地将视频中的帧分为2类:关键帧和非关键帧；模型在这2组框架上只是在Backbone阶段发生了变化。

通过这种方式可以在生成准确预测的同时保持快速运行时之间取得平衡。

3.3 部分功能转换

从邻近关键帧转换(即扭曲)特征被证明是一种有效的策略，以减少主干计算，以产生快速的视频Box目标检测器。具体地说，是使用离网光流网络转换所有主干特征。然而，由于光流估计中不可避免的误差，我们发现它不能提供像实例分割这样的像素级任务所需的足够精确的特征。在本研究中提出执行部分特征转换，以改善转换后的特征的品质，同时仍维持快速的运行时间。

总而言之，部分特征转换设计产生了更高质量的特征图，同时也支持实时速度。

3.4 有效的运动估计

该部分描述如何有效地计算关键帧和非关键帧之间的流。

为了进行快速的特征变换，需要有效地估计目标运动。现有的进行流导向特征变换的框架直接采用现成的像素级光流网络进行运动估计。例如，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

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相同的损失来训练模型:

对于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

4 实验

4.1 精度数据

4.2 实际效果

5 参考

[1].YolactEdge: Real-time Instance Segmentation on the Edge
[2].https://github.com/haotian-liu/yolact_edge

原文：[ 集智书童 ]
作者：ChaucerG
​

推荐阅读

• 项目实践 | 基于YOLO-V5实现行人社交距离风险提示
• 可以让深度学习编译器来指导算子优化吗
• MindSpore AI科学计算系列（5）：AI框架加速海洋数值模拟