Mask R-CNN 实例分割

Mask R-CNN 实例分割

本节摘要:给 Faster R-CNN 检测器加一个小小的掩码分支,你就有了实例分割。本节端到端走通 Mask R-CNN 的架构:骨干、FPN、RPN、RoIAlign、框头、掩码头,讲清难点所在——RoIAlign 为何取代 RoIPool(不量化的双线性采样,免费给掩码 AP 加 3-4 个点)。我们手写 RoIAlign 并与 torchvision 对齐验证,加载预训练 maskrcnn_resnet50_fpn_v2,读对它的输出格式,并通过换框头与掩码头、冻结骨干的方式在小数据集上微调。读完本节,你能用 torchvision 在自定义类别上跑出生产级实例掩码。

对应原课程:Phase 4 · Lesson 08 · instance-segmentation-mask-rcnn(原英文 phases/04-computer-vision/08-instance-segmentation-mask-rcnn/docs/en.md)。

学习目标

阅读完本节,你应当能够:

1. 端到端描绘 Mask R-CNN 架构:骨干、FPN、RPN、RoIAlign、框头、掩码头。
2. 从零实现 RoIAlign,并解释为什么 RoIPool 不再使用。
3. 用 torchvision 的 maskrcnn_resnet50_fpn_v2 预训练模型产出生产级实例掩码,并读对它的输出格式。
4. 通过替换框头与掩码头、冻结骨干,在小自定义数据集上微调 Mask R-CNN。

一、问题与直觉

语义分割给你每类一个掩码。实例分割给你每个物体一个掩码,哪怕两个物体同类。计数个体、跨帧跟踪、测量物体(墙上每块砖的框、显微图里每个细胞)都要求实例分割。

Mask R-CNN(He 等,2017)通过把实例分割重新表述为「检测加掩码」解决了这个问题。设计干净到接下来的五年里,几乎每一篇实例分割论文都是 Mask R-CNN 变体,而 torchvision 实现至今仍是中小数据集的生产默认。

难的工程问题是采样:如何从一个角点不对齐像素边界的提议框里,裁出固定尺寸的特征区域?弄错这一点,处处掉零点几个 mAP。RoIAlign 就是答案。

架构

要理解的五个零件:

1. 骨干 —— ImageNet 上训练的 ResNet-50 或 ResNet-101,产生步幅 4、8、16、32 的特征图层级。
2. FPN(特征金字塔网络) —— 自顶向下 + 横向连接,给每一级 C 个通道的语义丰富特征。检测查询与物体尺寸匹配的 FPN 层级。
3. RPN(区域提议网络) —— 一个小卷积头,在每个锚框位置预测「这里有没有物体?」和「怎么精修框?」。每图约产出 1000 个提议。
4. RoIAlign —— 从任意 FPN 层级的任意框上,采样出固定尺寸(如 7×7)的特征块。双线性采样,不量化。
5. 头 —— 两层框头精修框并选类,加一个小卷积头为每个提议输出 28×28 的二值掩码。

为什么是 RoIAlign 而不是 RoIPool

原始 Fast R-CNN 用 RoIPool:把提议框切成网格,每格取最大特征,所有坐标四舍五入到整数。这个舍入让特征图与输入像素坐标错开最多一整个特征图像素——在 224×224 图像上小,在步幅 32 的特征图上是灾难。

RoIPool:
  框 (34.7, 51.3, 98.2, 142.9)
  舍入 -> (34, 51, 98, 142)
  切网格 -> 舍入每个格子的边界
  每一步都累积错位

RoIAlign:
  框 (34.7, 51.3, 98.2, 142.9)
  在精确浮点坐标上用双线性插值采样
  处处不舍入
​

RoIAlign 在 COCO 上免费给掩码 AP 加 3-4 个点。每一个在乎定位的检测器现在都用它——YOLOv7 seg、RT-DETR、Mask2Former 都一样。

RPN 一段话讲清

在特征图的每个位置,放 K 个不同尺寸和形状的锚框。为每个锚框预测一个 objectness 分数和一个回归偏移,把锚框变成更贴合的框。按分数保留 top ~1000 个框,在 IoU 0.7 处做 NMS,把幸存者交给头。RPN 用它自己的迷你损失训练——结构与第 6 节的 YOLO 损失一样,只是两类(物体 / 非物体)。

掩码头

对每个提议(RoIAlign 之后),掩码头是一个微型 FCN:四个 3×3 卷积、一个 2× 反卷积、一个最终 1×1 卷积,在 28×28 分辨率上产出 num_classes 个输出通道。只保留与预测类别对应的通道,其余忽略。这把掩码预测与分类解耦。

把 28×28 掩码上采样到提议的原始像素尺寸,得到最终二值掩码。

损失

Mask R-CNN 有四个损失相加:

L = L_rpn_cls + L_rpn_box + L_box_cls + L_box_reg + L_mask
​

• L_rpn_cls, L_rpn_box —— RPN 提议的 objectness + 框回归。
• L_box_cls —— 头分类器上 (C+1) 类(含背景)的交叉熵。
• L_box_reg —— 头框精修的 smooth L1。
• L_mask —— 28×28 掩码输出上的逐像素二值交叉熵。

每个损失有自己的默认权重;torchvision 实现把它们作为构造参数暴露。

输出格式

torchvision.models.detection.maskrcnn_resnet50_fpn_v2 返回一个 dict 列表,每图一个:

{
    "boxes":  (N, 4), (x1, y1, x2, y2) 像素坐标,
    "labels": (N,) 类 ID,0 = 背景所以下标从 1 开始,
    "scores": (N,) 置信度,
    "masks":  (N, 1, H, W) [0, 1] 浮点掩码,阈值 0.5 得二值,
}
​

掩码已是全图分辨率。28×28 头输出已在内部上采样。

二、从零实现

步骤 1:从零实现 RoIAlign

这是 Mask R-CNN 里一个用代码比用文字更好理解的零件。

import torch
import torch.nn.functional as F

def roi_align_single(feature, box, output_size=7, spatial_scale=1 / 16.0):
    """
    feature: (C, H, W) 单图特征图
    box: (x1, y1, x2, y2) 原始图像像素坐标
    output_size: 输出网格边长(框头 7,掩码头 14)
    spatial_scale: 特征图步幅的倒数
    """
    C, H, W = feature.shape
    x1, y1, x2, y2 = [c * spatial_scale - 0.5 for c in box]
    bin_w = (x2 - x1) / output_size
    bin_h = (y2 - y1) / output_size

    grid_y = torch.linspace(y1 + bin_h / 2, y2 - bin_h / 2, output_size)
    grid_x = torch.linspace(x1 + bin_w / 2, x2 - bin_w / 2, output_size)
    yy, xx = torch.meshgrid(grid_y, grid_x, indexing="ij")

    gx = 2 * (xx + 0.5) / W - 1
    gy = 2 * (yy + 0.5) / H - 1
    grid = torch.stack([gx, gy], dim=-1).unsqueeze(0)
    sampled = F.grid_sample(feature.unsqueeze(0), grid, mode="bilinear",
                            align_corners=False)
    return sampled.squeeze(0)
​

每个数都在双线性采样位置上,没有舍入、没有量化、没有丢失梯度。

步骤 2:与 torchvision 的 RoIAlign 比较

from torchvision.ops import roi_align

feature = torch.randn(1, 16, 50, 50)
boxes = torch.tensor([[0, 10, 20, 100, 90]], dtype=torch.float32)  # (batch_idx, x1, y1, x2, y2)

ours = roi_align_single(feature[0], boxes[0, 1:].tolist(), output_size=7, spatial_scale=1/4)
theirs = roi_align(feature, boxes, output_size=(7, 7), spatial_scale=1/4, sampling_ratio=1, aligned=True)[0]

print(f"shape ours:   {tuple(ours.shape)}")
print(f"shape theirs: {tuple(theirs.shape)}")
print(f"max|diff|:    {(ours - theirs).abs().max().item():.3e}")
​

sampling_ratio=1、aligned=True 时,两者在 1e-5 内匹配。

步骤 3:加载预训练 Mask R-CNN

import torch
from torchvision.models.detection import maskrcnn_resnet50_fpn_v2, MaskRCNN_ResNet50_FPN_V2_Weights

model = maskrcnn_resnet50_fpn_v2(weights=MaskRCNN_ResNet50_FPN_V2_Weights.DEFAULT)
model.eval()
print(f"params: {sum(p.numel() for p in model.parameters()):,}")
print(f"classes (including background): {len(model.roi_heads.box_predictor.cls_score.out_features * [0])}")
​

4600 万参数,91 类(COCO)。第一类(id 0)是背景;模型实际检测的一切从 id 1 开始。

步骤 4:跑推理

with torch.no_grad():
    x = torch.randn(3, 400, 600)
    predictions = model([x])
p = predictions[0]
print(f"boxes:  {tuple(p['boxes'].shape)}")
print(f"labels: {tuple(p['labels'].shape)}")
print(f"scores: {tuple(p['scores'].shape)}")
print(f"masks:  {tuple(p['masks'].shape)}")
​

掩码张量形状 (N, 1, H, W)。阈值 0.5 得每个物体的二值掩码:

binary_masks = (p['masks'] > 0.5).squeeze(1)  # (N, H, W) 布尔
​

步骤 5:为自定义类别数换头

常见微调配方:复用骨干、FPN、RPN;替换两个分类头。

from torchvision.models.detection.faster_rcnn import FastRCNNPredictor
from torchvision.models.detection.mask_rcnn import MaskRCNNPredictor

def build_custom_maskrcnn(num_classes):
    model = maskrcnn_resnet50_fpn_v2(weights=MaskRCNN_ResNet50_FPN_V2_Weights.DEFAULT)
    in_features = model.roi_heads.box_predictor.cls_score.in_features
    model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)
    in_features_mask = model.roi_heads.mask_predictor.conv5_mask.in_channels
    hidden_layer = 256
    model.roi_heads.mask_predictor = MaskRCNNPredictor(in_features_mask, hidden_layer, num_classes)
    return model

custom = build_custom_maskrcnn(num_classes=5)
print(f"custom cls_score.out_features: {custom.roi_heads.box_predictor.cls_score.out_features}")
​

num_classes 必须含背景类,所以 4 个物体类的数据集用 num_classes=5。

步骤 6:冻结不需要训练的部分

小数据集上,冻结骨干和 FPN。只有 RPN 的 objectness + 回归和两个头在学习。

def freeze_backbone_and_fpn(model):
    # torchvision 的 Mask R-CNN 把 FPN 打包在 `model.backbone` 里
    # (作为 `model.backbone.fpn`),所以遍历 `model.backbone.parameters()`
    # 同时覆盖 ResNet 特征层和 FPN 横向/输出卷积。
    for p in model.backbone.parameters():
        p.requires_grad = False
    return model

custom = freeze_backbone_and_fpn(custom)
trainable = sum(p.numel() for p in custom.parameters() if p.requires_grad)
print(f"trainable after freeze: {trainable:,}")
​

在 500 张图的数据集上,这是收敛与过拟合的分水岭。

三、框架对比

torchvision 里 Mask R-CNN 的完整训练循环 40 行,任务之间没有本质变化——换数据集就跑。

def train_step(model, images, targets, optimizer):
    model.train()
    loss_dict = model(images, targets)
    losses = sum(loss for loss in loss_dict.values())
    optimizer.zero_grad()
    losses.backward()
    optimizer.step()
    return {k: v.item() for k, v in loss_dict.items()}
​

targets 列表必须含每图 dict,带 boxes、labels、masks(作为 (num_instances, H, W) 二值张量)。训练时模型返回四损失的 dict,评估时返回预测列表,按 model.training 切换。

pycocotools 评估器对框和掩码都产出 mAP@IoU=0.5:0.95;两个数都要,才知道是框头还是掩码头是瓶颈。

四、可复用产物

本节产出两个可复用文件(位于原课程 outputs/):

• prompt-instance-vs-semantic-router.md:一个提示词——问三个问题,在实例/语义/全景分割之间挑,并指出该从哪个模型开始。
• skill-mask-rcnn-head-swapper.md:一个技能——给定新的 num_classes,为任意 torchvision 检测模型生成换头的那 10 行代码。

五、练习

1. (简单) 在 100 个随机框上把你的 RoIAlign 对 torchvision.ops.roi_align 验证。报告最大绝对差。再跑 RoIPool(2017 前的行为),展示它在靠近边界的框上偏移约 1-2 个特征图像素。
2. (中等) 在一个 50 张图的自定义数据集(任意两类:气球、鱼、坑洞、logo)上微调 maskrcnn_resnet50_fpn_v2。冻结骨干,训 20 个 epoch,报告掩码 AP@0.5。
3. (困难) 把 Mask R-CNN 的掩码头换成在 56×56 而非 28×28 上预测的。测 mAP@IoU=0.75 前后。解释增益(或没有)为何符合预期的边界精度 / 显存权衡。

本节要点回顾

1. 实例分割 = 检测 + 掩码——Mask R-CNN 给 Faster R-CNN 加一个 FCN 头,每提议每类产 28×28 二值掩码。
2. 五大零件——骨干、FPN(每级 C 通道语义特征)、RPN(约 1000 个提议)、RoIAlign、框头+掩码头。
3. RoIAlign 是关键——双线性采样固定尺寸特征网格,不量化;比 RoIPool 免费加 3-4 个掩码 AP。
4. RPN 与 YOLO 损失同构——只是两类(物体/非物体),按分数 NMS 后交给头。
5. 掩码头解耦分类——每类一个掩码通道,只保留预测类。
6. 四个损失——RPN 的 cls+box、头的 cls+reg、掩码的逐像素 BCE。
7. 输出格式——boxes/labels/scores/masks,0 是背景、下标从 1 开始,掩码已是全图分辨率。
8. 微调配方——复用骨干+FPN+RPN,换 FastRCNNPredictor 与 MaskRCNNPredictor,num_classes 含背景。
9. 小数据冻结骨干+FPN——只训 RPN 与两个头,500 张图上是收敛与过拟合的分水岭。
10. 同时报告框与掩码 mAP——才知道是哪个头是瓶颈。

下一节,我们离开「给真实图像打标签」,进入「从噪声合成图像」——用 GAN 讲清生成对抗的博弈,以及为什么它会模式崩溃。