使用mmdetection2.0实现SOLOV2-全新的实例分割框架

前端工程师，Flutter开发，AI，浙大硕士毕业

使用mmdetection2.0实现SOLOV2-全新的实例分割框架

本文由林大佬原创，转载请注明出处，来自腾讯、阿里等一线AI算法工程师组成的QQ交流群欢迎你的加入: 1037662480

笔者近日在做一些实例分割到tensorrt的操作，发现，对于许多二阶段的实例分割算法，转换到另一个推理平台很麻烦，主要是这里面存在RPN的一些操作，即便是类似于CenterMask这样的模型，也需要先获取box，然后经过ROIAlign和ROIPool来得到需要的mask，这个过程不仅计算繁琐，而且很难导出到我们想要的模型格式，比如ONNX就不支持这里面的很多op。当我看到这张图的时候，我发现，SOLOV2的性能似乎比maskrcnn差不多，推理速度只要一半，并且速度和精确度都超过了BlendMask。今天我们就用mmdetection 2.0的版本来实现一个SOLOV2.

就目前来讲，很多实例分割算法存在的问题弊端主要是：

速度太慢了，比如Maskrcnn，虽然声名远扬，但是要把它部署到realtime，还是很困难的；
精度不够，比如Yolact，即便是Yolact++，其精度其实也只能说差强人意，连MaskRCNN都比不上的实力分割，速度再快，也会限制它的使用场景；
BlendMask，CenterMask这类的算法，都差不多，基于FCOS构建，本质上没啥区别，还是和MaskRCNN整体流程差不多，只不过检测器变了而已，对于部署来说依旧很麻烦。

当然本文要解决的问题并不是部署问题，而是告诉大家，我们有一个速度更快，精度更好的模型，而且部署相对来说可能更好一点。那么第一步是不是得用python实现一波呢？

SOLOV2

对于SOLO系列算法，具体的来龙去脉就不赘述了。先放一个指标对比：

对于V2来讲，变化最大的可能就是这个动态的mask head：

class SOLOV2Head:
    ....
    def forward(self, feats, eval=False):
            new_feats = self.split_feats(feats)
            featmap_sizes = [featmap.size()[-2:] for featmap in new_feats]
            upsampled_size = (feats[0].shape[-2], feats[0].shape[-3])
            kernel_pred, cate_pred = multi_apply(self.forward_single, new_feats,
                                              list(range(len(self.seg_num_grids))),
                                              eval=eval)
            # add coord for p5
            x_range = torch.linspace(-1, 1, feats[-2].shape[-1], device=feats[-2].device)
            y_range = torch.linspace(-1, 1, feats[-2].shape[-2], device=feats[-2].device)
            y, x = torch.meshgrid(y_range, x_range)
            y = y.expand([feats[-2].shape[0], 1, -1, -1])
            x = x.expand([feats[-2].shape[0], 1, -1, -1])
            coord_feat = torch.cat([x, y], 1)
            feature_add_all_level = self.feature_convs[0](feats[0]) 
            for i in range(1,3):
                feature_add_all_level = feature_add_all_level + self.feature_convs[i](feats[i])
            feature_add_all_level = feature_add_all_level + self.feature_convs[3](torch.cat([feats[3],coord_feat],1))

            feature_pred = self.solo_mask(feature_add_all_level)   
            N, c, h, w = feature_pred.shape
            feature_pred = feature_pred.view(-1, h, w).unsqueeze(0)
            ins_pred = []

            for i in range(5):
                kernel = kernel_pred[i].permute(0,2,3,1).contiguous().view(-1,c).unsqueeze(-1).unsqueeze(-1)
                ins_i = F.conv2d(feature_pred, kernel, groups=N).view(N,self.seg_num_grids[i]**2, h,w)
                if not eval:
                    ins_i = F.interpolate(ins_i, size=(featmap_sizes[i][0]*2,featmap_sizes[i][1]*2), mode='bilinear')
                if eval:
                    ins_i=ins_i.sigmoid()
                ins_pred.append(ins_i)
            return ins_pred, cate_pred