目录
OrienMask是2021年新出的一篇实时的实例分割模型。下图是OrienMask在coco数据集上的测试结果。在单张2080ti上可以达到42.7fps的速度。
这篇文章的主体思路是先使用yolo检测出目标的bbox，然后再检测出bbox中的前景来完成实例分割任务。
网络结构
上图是OrienMask的网络结构，主要分为3块：Backbone、FPN、Heads。
使用darknet53作为Backbone，进行了32倍下采样，分别在4x，8x，16x，32x下采样的进行FPN来进行特征融合；最后是检测heads，32x、16x、8倍下采样分支上输出bbox结果，4倍下采样的分支上输出Orientation Maps。可以看出来除了在4倍下采样的分支上输出Orientation Maps之外只是一个yolo。
网络输出
这里详细介绍网络的输出和Orientation Maps的作用。
在后处理之前，模型的输出是：dets_conf，dets_coor，dets_orien。
其中dets_conf的尺寸是：18207*80。这里 18207 = 3 ∗ ( 1 7 2 + 3 4 2 + 6 8 2 ) 18207=3*(17^2+34^2+68^2) 18207=3∗(172+342+682)，3是表示每个尺度下有3个不同比例的anchor，17、34、68分别是3个尺度的grid size；80代表类的数量。dets_conf里存放的是所有尺度的grid的置信度。
dets_coor的尺寸是：18207*4。18207的定义同上，4代表bbox的长度即xywh。这里存放的所有尺度的grid的xywh的偏移量。
dets_orien的尺寸是9*2*544*544。这里9代表9个anchor（3个尺度*3个不同长宽比），2代表x,y（向量），544是图像输入尺寸。即dets_orien里存了输入图像在9个anchor上的方向向量。
后处理
怕表达不清楚，后处理这块照着代码聊吧。这里提供的代码是我转tensorrt时没能转进onnx的部分，与源码中的后处理代码可能会有一点出入，但是最终结果是对的。后面会介绍转tensorrt的全部流程和代码。
def post_processing(dets_coord, dets_conf, dets_orien, dets_anchor_idx, grid_sizes, base_xy, grid_anchors, pre_post_num=400, conf_thresh=0.005):
selected_inds, dets_cls = torch.nonzero(dets_conf > conf_thresh, as_tuple=True)
selected_inds = selected_inds.view(-1)
dets_cls = dets_cls.view(-1)
dets_conf = dets_conf[selected_inds, dets_cls]
# 选择固定数量的候选框留给 NMS 操作
if selected_inds.numel() > pre_post_num:
dets_conf, topk_inds = dets_conf.topk(pre_post_num)
selected_inds = selected_inds[topk_inds]
dets_cls = dets_cls[topk_inds]
dets_coord = dets_coord[selected_inds]
dets_anchor_idx = dets_anchor_idx[selected_inds]
dets_orien = dets_orien * grid_anchors.reshape(3, 6, 1, 1) / 2
final_pred = multi_class_nms(base_xy, grid_sizes, dets_coord, dets_conf, dets_cls, dets_anchor_idx, dets_orien)
return final_pred
第2到5行是选择出置信度大于阈值的，selected_inds是anchor的索引，dets_cls是类别索引，这里的阈值设的非常小，会影响速率，这个问题在讲nms部分时会分析。第7行的if是防止太多的数据进入nms，即满足置信度条件的数量如果大于pre_post_num就取个topk。直到第15行，进入multi_class_nms，这里不仅做了bb的nms还根据dets_orien求出mask图。
下面展开multi_class_nms代码：
def batched_nms(dets, cats, threshold=0.5, normalized=True):
"""Batched non maximum suppression
Refer to `torchvision.ops.boxes.batched_nms` for details.
Args:
dets (n, 6): bounding boxes with x, y, width, height and score
cats (n,): cls indices of bounding boxes
threshold (float): IoU threshold for NMS
normalized (bool): use normalized coordinates or not
Returns:
remains dets, cats and their indices
"""
if dets.size(0) == 0:
keep = dets.new_zeros(0, dtype=torch.long)
else:
max_coordinate = 1.5 if normalized else dets[:, :2].max() + dets[:, 2:4].max() / 2
offsets = cats.float().view(-1, 1) * (max_coordinate + 0.5)
dets_for_nms = dets.clone()
dets_for_nms[:, :2] += offsets
print("dets_for_nms: ", dets_for_nms.shape)
print(dets_for_nms)
print("*"*88)
if dets.is_cuda:
keep = nms_cuda.nms(dets_for_nms, threshold)
else:
keep = nms_cpu.nms(dets_for_nms, threshold)
return dets[keep], cats[keep], keep
def multi_class_nms(base_xy, grid_sizes, pred_coord, pred_conf, pred_cls, pred_anchor_idx, pred_orien, post_nms_num=100, orien_thresh=0.3):
dets = torch.cat([pred_coord, pred_conf.unsqueeze(-1)], dim=1)
# 使用Batch_NMS求取结果
dets, cats, keep = batched_nms(dets, pred_cls)
# 求取前100个结果
if keep.numel() > post_nms_num:
_, topk_inds = dets[:, -1].topk(post_nms_num)
dets = dets[topk_inds]
cats = cats[topk_inds]
keep = keep[topk_inds]
anchor_idx = pred_anchor_idx[keep]
x_centers = (grid_sizes[anchor_idx, 0] * dets[:, 0]).view(-1, 1, 1)
y_centers = (grid_sizes[anchor_idx, 1] * dets[:, 1]).view(-1, 1, 1)
det_width = dets[:, 2].view(-1, 1, 1)
det_height = dets[:, 3].view(-1, 1, 1)
pred_orien = pred_orien.reshape(9, 2, 544, 544)
pred_orien = pred_orien + base_xy
masks = ((torch.abs(pred_orien[anchor_idx, 0] - x_centers) < orien_thresh * det_width * grid_sizes[anchor_idx, 0].view(-1, 1, 1)) &
(torch.abs(pred_orien[anchor_idx, 1] - y_centers) < orien_thresh * det_height * grid_sizes[anchor_idx, 1].view(-1, 1, 1)))
return {'bbox': dets, 'mask': masks, 'cls': cats}
multi_class_nms最开始将xywh偏移量（pred_coord）和置信度（pred_conf）cat到一起送到batched_nms里。batched_nms里主要做了bbox的nms，在这里又一次卡了置信的的阈值，这次阈值比较大也是最终的bbox的阈值，个人感觉这样做是不太合理的，卡一次最终的阈值就可以了，因为在C++复现后处理的时候会有速率的影响，并且我在C++那里把卡阈值的操作做成了函数，连续调用两次相同的函数看起来就很呆。batched_nms是非常常用的东西就不展开介绍了，重点是根据pred_orien生成mask图：由于预测出的orien是单位向量，所以在post_processing中先乘上了grid_anchors，再在multi_class_nms里加上了base_xy。
这里grid_anchors是anchor的尺寸，在这个模型里设置的是：
grid_anchors = torch.tensor([[1.5000, 2.0000],
[2.3750, 4.5000],
[5.0000, 3.5000],
[2.2500, 4.6875],
[4.7500, 3.4375],
[4.5000, 9.1250],
[4.4375, 3.4375],
[6.0000, 7.5938],
[14.3438, 12.5312]])/2
base_xy是在不同尺度grid上的伪像素坐标，生成方式我也摘出来了：
grid_size = [[17, 17],
[34, 34],
[68, 68]]
nHs = [size[0] for size in grid_size]
nWs = [size[1] for size in grid_size]
num_anchors = [len(m) for m in anchor_mask]
base_xy = torch.zeros(9, 2, 544, 544)
grid_x, grid_y, anchor_idx = [], [], []
for m, nA, nH, nW in zip(anchor_mask, num_anchors, nHs, nWs):
# 构建一个基础的方向坐标 (unit size=grid)
base_y, base_x = torch.meshgrid([torch.arange(544, dtype=torch.float) / 544 * nH,
torch.arange(544, dtype=torch.float) / 544 * nW])
# 第一个维度上进行拼接操作
base_xy[m] = torch.stack([base_x, base_y], dim=0)
最后orien向量的x<bbox.width*threshold 且 y<bbox.height*threshold时认为时bbox的前景。这样就完成了实例分割。
总结
这个模型的训练工作是其他同事来做的，我的主要工作是把这个模型移植到tensorrt上所以更多关注了模型的后处理部分，模型架构部分的代码没去细看。整体看下来这个模型架构比较清晰，后处理过程也比较简单。我的初版代码在tensorrt上infer时间是18ms后处理大部分放到了cuda上是3ms，预处理使用opencv写的会比较耗时大概20ms，但是预处理可以和infer+后处理放在不同线程去跑。综合帧率还是能达到40fps左右的（速度上没体现出来tensorrt的优势，惭愧）。代码优化后并且把模型转fp16或者量化到int8应该会更快的。后续有时间的话再记录转tensorrt的过程。