CondInst：性能和速度均超越Mask RCNN的實(shí)例分割模型

點(diǎn)藍(lán)色字關(guān)注“機(jī)器學(xué)習(xí)算法工程師”

設(shè)為星標(biāo)，干貨直達(dá)！

對于實(shí)例分割來說，主流的做法還是基于先檢測后分割的流程，比如最流行的Mask RCNN模型就是構(gòu)建在Faster RCNN基礎(chǔ)上。目前基于one-stage的物體檢測模型已經(jīng)在速度和性能上超越two-stage模型，同樣地，大家也希望能找到one-stage的實(shí)例分割模型來替換Mask RCNN。目前這方面的工作主要集中在三個(gè)方向：

• Mask encoding：對2D mask編碼為1D representation，比如PolarMask基于輪廓構(gòu)建了polar representation，而MEInst則將mask壓縮成一個(gè)1D vector，這樣預(yù)測mask就類似于box regress那樣直接加在one-stage檢測模型上；
• 分離檢測和分割：將檢測和分割分離成兩個(gè)部分這樣可以并行化，如YOLACT在檢測模型基礎(chǔ)上額外預(yù)測了一系列prototype masks，然后檢測部分每個(gè)instance會預(yù)測mask coeffs來組合masks來產(chǎn)生instance mask，BlendMask是對這一工作的進(jìn)一步改進(jìn)；
• 不依賴檢測的實(shí)例分割：不依賴檢測框架直接進(jìn)行實(shí)例分割，TensorMask和SOLO屬于此種類型，前者速度太慢，后者速度和效果都非常好；

對于mask encoding方法，雖然實(shí)現(xiàn)起來比較容易，但是往往會造成2D mask的細(xì)節(jié)損失，所以性能上會差一點(diǎn)；分離檢測和分割，對于分割部分可以像語義分割那樣預(yù)測global mask，分辨率上會更高（要知道Mask RCNN的mask分辨率只有28x28），但是這種方法需要一種好的方式來產(chǎn)生instance mask；不依賴檢測而直接進(jìn)行實(shí)例分割這可能是未來的趨勢。這里介紹的CondInst，其實(shí)屬于第二種，但是它與YOLACT不同，其核心點(diǎn)是檢測部分為每個(gè)instance預(yù)測不同的mask head，然后基于global mask features來產(chǎn)生instance mask，思路非常簡單，而且實(shí)現(xiàn)起來也極其容易（已經(jīng)開源在AdelaiDet），更重要的是速度和效果上均超越Mask RCNN。

整體結(jié)構(gòu)

CondInst是構(gòu)建在物體檢測模型FCOS之上的（CondInst和FCOS是同一個(gè)作者），所以理解CondInst必須先理解FCOS，可以參考之前關(guān)于FCOS的介紹文章（FCOS），但其實(shí)CondInst也可以依賴其他的one-stage模型，CondInst整體結(jié)構(gòu)如下圖所示：

相比FCOS，CondInst多了一個(gè)mask branch，其得到的mask features將作為mask FCN的輸入來生成最終的instance mask，這個(gè)mask features來自于P3，所以大小是輸入圖像的1/8。另外在FCOS的檢測部分增加了controller head（實(shí)際上controller head是直接加在box head上的），用來產(chǎn)生每個(gè)instance的mask head網(wǎng)絡(luò)的參數(shù)。這個(gè)思想其實(shí)是CondConv，傳統(tǒng)的Conv訓(xùn)練完成后是固定的filters，但是CondConv的filters基于input和一個(gè)另外的網(wǎng)路來動態(tài)產(chǎn)生的。CondInst用來controller head生成instance-aware的mask FCN head，每個(gè)instance都有自己獨(dú)有的mask head，instance的形狀和大小等信息都編碼在其中。所以當(dāng)mask head作用在global mask features上時(shí)，就可以區(qū)分當(dāng)前的instance和其它背景信息，從而預(yù)測出instance mask。

這樣CondInst就可以實(shí)現(xiàn)實(shí)例分割了，CondInst的正負(fù)樣本策略和FCOS一樣，都是通過center region sampling方式來決定正負(fù)樣本，其訓(xùn)練的loss相比FCOS增加intance mask的loss，這個(gè)loss也只計(jì)算正樣本部分：

Mask Branch

CondInst的mask branch就和語義分割類似是一個(gè)FCN網(wǎng)絡(luò)，包括4個(gè)channel為128的3x3卷積，然后最后接一個(gè)channel為8的1x1卷積。mask branch輸入為FPN的P3特征，所以最終產(chǎn)生的特征為原始輸入圖像的1/8，特征channel為8，之所以用一個(gè)較小的channel是為了減少controller head所需生成的參數(shù)量，而且實(shí)驗(yàn)中發(fā)現(xiàn)采用較小的channel就夠了，當(dāng)channel為2時(shí)mask AP僅掉了0.3%。不過從開源的代碼來看，mask branch的輸入應(yīng)該是來自于FPN的P3，P4和P5，具體實(shí)現(xiàn)上先將P4和P5的特征通過雙線性插值，然后和P3加到一起作為mask branch的輸入。就像YOLACT一樣，mask branch產(chǎn)生的特征還可以額外加上語義分割的loss來進(jìn)行輔助，這個(gè)不會影響inference過程，但是實(shí)驗(yàn)上mask AP大約可以提升1個(gè)點(diǎn)，具體實(shí)現(xiàn)上如下：

#?額外的語義loss，采用focal?loss
if?self.training?and?self.sem_loss_on:
????logits_pred?=?self.logits(self.seg_head(
????????????????features[self.in_features[0]]
????????????))?#?預(yù)測logits，區(qū)分class

????#?計(jì)算語義分割的gt，這里的原則是合并instance的gt?mask，但是當(dāng)不同instance有重疊時(shí)，會取面積最小的instance的class作為gt
????semantic_targets?=?[]
????for?per_im_gt?in?gt_instances:
????????h,?w?=?per_im_gt.gt_bitmasks_full.size()[-2:]
????????areas?=?per_im_gt.gt_bitmasks_full.sum(dim=-1).sum(dim=-1)
????????areas?=?areas[:,?None,?None].repeat(1,?h,?w)
????????areas[per_im_gt.gt_bitmasks_full?==?0]?=?INF
????????areas?=?areas.permute(1,?2,?0).reshape(h?*?w,?-1)
????????min_areas,?inds?=?areas.min(dim=1)
????????per_im_sematic_targets?=?per_im_gt.gt_classes[inds]?+?1
????????per_im_sematic_targets[min_areas?==?INF]?=?0
????????per_im_sematic_targets=?per_im_sematic_targets.reshape(h,?w)
????????semantic_targets.append(per_im_sematic_targets)

????semantic_targets?=?torch.stack(semantic_targets,?dim=0)?#?[N,?1,?H,?W]?

????#?對gt進(jìn)行降采樣，為原始的1/8
????semantic_targets?=?semantic_targets[:,?None,?self.out_stride?//?2::self.out_stride,?self.out_stride?//?2::self.out_stride]

????#?one-hot?gt
????num_classes?=?logits_pred.size(1)
????class_range?=?torch.arange(num_classes,?dtype=logits_pred.dtype,?device=logits_pred.device)[None,?:,?None,?None]
????class_range?=?class_range?+?1
????one_hot?=?(semantic_targets?==?class_range).float()
????num_pos?=?(one_hot?>?0).sum().float().clamp(min=1.0)
????
????#?采用focal?loss
????loss_sem?=?sigmoid_focal_loss_jit(
????????????????logits_pred,?one_hot,
????????????????alpha=self.focal_loss_alpha,
????????????????gamma=self.focal_loss_gamma,
????????????????reduction="sum",
????????????)?/?num_pos
????????????losses['loss_sem']?=?loss_sem

return?mask_feats,?losses

Controller Head

前面說過，CondInst的核心就在于controller head，其用來產(chǎn)生mask head的網(wǎng)絡(luò)參數(shù)，這個(gè)參數(shù)是每個(gè)instance所獨(dú)有的，所以當(dāng)輸入為全局mask特征時(shí)，可以預(yù)測出instance mask。由于controller head會編碼instance的形狀和大小信息，所以它是直接加在FCOS的box head上的，就和centerness head一樣。

controller head的輸出channel數(shù)為N，恰好是mask head的網(wǎng)絡(luò)參數(shù)量。mask head采用一個(gè)輕量級的FCN網(wǎng)絡(luò)，包含三個(gè)channel為8的3x3卷積層，卷積之后接ReLU，最后一層卷積直接加上sigmoid（二分類）就可以預(yù)測instance mask。所以mask head的參數(shù)量N為169：(#weights = (8 + 2) × 8(conv1) + 8 × 8(conv2) + 8 × 1(conv3) and #biases = 8(conv1) + 8(conv2) + 1(conv3))。這里的輸入channel是8+2，而不是8，是因?yàn)樗腿雖ask head的輸入除了包括，還包含relative coordinates maps，即相對于當(dāng)前instance的位置(x,y)的相對位置坐標(biāo)，在實(shí)現(xiàn)上只需要把x,y的relative coordinates maps與拼接在一起即可，如果去掉相對位置maps，CondInst性能下降比較厲害，其實(shí)也合理，因?yàn)閏ontroller head產(chǎn)生的mask head參數(shù)是由CNN得到的，它雖然可以編碼instance的shape信息，但是不能準(zhǔn)確地學(xué)習(xí)到instance在圖像中的位置，所以加上relative coordinates maps對準(zhǔn)確地分割當(dāng)前instance比較重要。另外論文中的一個(gè)有趣的實(shí)驗(yàn)是mask head只輸入relative coordinates maps就可以得到31.3%的mask AP值，這或許說明controller head產(chǎn)生的mask head足夠強(qiáng)大。下面是對controller head的輸出進(jìn)行結(jié)構(gòu)化解析的代碼：

def?parse_dynamic_params(params,?channels,?weight_nums,?bias_nums):
????assert?params.dim()?==?2
????assert?len(weight_nums)?==?len(bias_nums)
????assert?params.size(1)?==?sum(weight_nums)?+?sum(bias_nums)

????num_insts?=?params.size(0)
????num_layers?=?len(weight_nums)

????params_splits?=?list(torch.split_with_sizes(
????????params,?weight_nums?+?bias_nums,?dim=1
????))

????weight_splits?=?params_splits[:num_layers]
????bias_splits?=?params_splits[num_layers:]

????for?l?in?range(num_layers):
????????if?l?1:
????????????#?out_channels?x?in_channels?x?1?x?1
????????????weight_splits[l]?=?weight_splits[l].reshape(num_insts?*?channels,?-1,?1,?1)
????????????bias_splits[l]?=?bias_splits[l].reshape(num_insts?*?channels)
????????else:
????????????#?out_channels?x?in_channels?x?1?x?1
????????????weight_splits[l]?=?weight_splits[l].reshape(num_insts?*?1,?-1,?1,?1)
????????????bias_splits[l]?=?bias_splits[l].reshape(num_insts)

????return?weight_splits,?bias_splits

由于輸入的是1/8圖片大小，所以產(chǎn)生的instance mask也是1/8圖片大小，為了產(chǎn)生高質(zhì)量的mask，采用雙線性插值對預(yù)測的instance mask上采樣4x，那么最終輸出的instance mask就是1/2圖片大小。mask部分的loss采用的是dice loss，它和focal loss一樣可以解決正負(fù)樣本不均衡問題，在計(jì)算時(shí)將gt mask下采樣2x以和預(yù)測mask達(dá)到同樣的大小。

Inference

CondInst的inference就比較直接了，首先是檢測部分得到檢測的結(jié)果，然后采用box-based NMS來去除重復(fù)框，最后選出top 100的檢測框，只有這部分instances會進(jìn)行instance mask的預(yù)測。由于產(chǎn)生的mask head非常小，所以100個(gè)instance的mask預(yù)測時(shí)間只需要4.5ms，那么CondInst的預(yù)測時(shí)間僅比原始的FCOS增加了約10%。這里額外要說的一點(diǎn)是CondInst的box預(yù)測主要用于NMS，但不會參與instance mask的預(yù)測中，而Mask R-CNN是需要box來進(jìn)行ROI croping。CondInst和其它實(shí)例分割在COCO上的效果對比如下：

此外，CondInst的作者近期又發(fā)布了一篇新的不錯(cuò)的工作：BoxInst，只用box級別的標(biāo)注就可以訓(xùn)練出一個(gè)不錯(cuò)的實(shí)例分割模型，這個(gè)模型也是構(gòu)建在CondInst上，只不過設(shè)計(jì)了兩個(gè)新的loss來進(jìn)行半監(jiān)督式的訓(xùn)練。最后放一個(gè)BoxInst的一個(gè)分割視頻demo：

參考

1. Conditional Convolutions for Instance Segmentation
2. AdelaiDet
3. ? ?BoxInst: High-Performance Instance Segmentation with Box Annotations