Mask R-CNN

Setup#

pip install torch torchvision

Everything below runs on CPU and downloads nothing.

The rounding bug, made visible#

Start with the part that matters. roi_pool rounds box coordinates onto the feature grid; roi_align samples at exact positions with bilinear interpolation.

The test is to move a box by half a pixel and see whether the output notices.

import torch, torch.nn.functional as F
from torchvision.ops import roi_align, roi_pool

# a feature map whose values encode their own column, so misalignment is visible
feat = torch.arange(64, dtype=torch.float32).reshape(1, 1, 8, 8) % 8
print("feature map (value = column index):")
print(feat[0, 0].numpy())

box   = torch.tensor([[0., 1.5, 1.5, 5.5, 5.5]])     # [batch_idx, x1, y1, x2, y2]
shift = torch.tensor([[0., 2.0, 1.5, 6.0, 5.5]])     # same box, moved half a pixel right

print("\nRoIPool, box x1=1.5 :\n", roi_pool(feat, box, output_size=2, spatial_scale=1.0)[0, 0].numpy())
print("RoIPool, box x1=2.0 :\n", roi_pool(feat, shift, output_size=2, spatial_scale=1.0)[0, 0].numpy())
print("-> identical. RoIPool rounded 1.5 and 2.0 onto the same grid cell.")

print("\nRoIAlign, box x1=1.5 :\n", roi_align(feat, box, output_size=2, spatial_scale=1.0, sampling_ratio=2)[0, 0].numpy())
print("RoIAlign, box x1=2.0 :\n", roi_align(feat, shift, output_size=2, spatial_scale=1.0, sampling_ratio=2)[0, 0].numpy())
print("-> the half-pixel shift moved the numbers. No rounding happened.")

# ---- pasting a 28x28 mask back into the image ----
m = torch.zeros(1, 1, 28, 28); m[0, 0, 6:22, 6:22] = 3.0        # logits from the mask head
x1, y1, x2, y2 = 10.4, 20.7, 46.9, 60.2
w, h = round(x2 - x1), round(y2 - y1)
resized = F.interpolate(m, size=(h, w), mode="bilinear", align_corners=False)
canvas = torch.zeros(1, 1, 80, 80)
canvas[0, 0, round(y1):round(y1) + h, round(x1):round(x1) + w] = (resized[0, 0] > 0).float()
print(f"\nmask head output {tuple(m.shape)} -> stretched to the box {(h, w)} -> pasted into {tuple(canvas.shape[-2:])}")
print("pixels on:", int(canvas.sum()), " rows used:", (canvas[0,0].sum(1) > 0).nonzero().flatten()[[0, -1]].tolist())

feature map (value = column index):
[[0. 1. 2. 3. 4. 5. 6. 7.]
 [0. 1. 2. 3. 4. 5. 6. 7.]
 [0. 1. 2. 3. 4. 5. 6. 7.]
 [0. 1. 2. 3. 4. 5. 6. 7.]
 [0. 1. 2. 3. 4. 5. 6. 7.]
 [0. 1. 2. 3. 4. 5. 6. 7.]
 [0. 1. 2. 3. 4. 5. 6. 7.]
 [0. 1. 2. 3. 4. 5. 6. 7.]]

RoIPool, box x1=1.5 :
 [[4. 6.]
 [4. 6.]]
RoIPool, box x1=2.0 :
 [[4. 6.]
 [4. 6.]]
-> identical. RoIPool rounded 1.5 and 2.0 onto the same grid cell.

RoIAlign, box x1=1.5 :
 [[2.5 4.5]
 [2.5 4.5]]
RoIAlign, box x1=2.0 :
 [[3. 5.]
 [3. 5.]]
-> the half-pixel shift moved the numbers. No rounding happened.

mask head output (1, 1, 28, 28) -> stretched to the box (40, 36) -> pasted into (80, 80)
pixels on: 528  rows used: [29, 52]

Reading that carefully#

Two different boxes, one identical RoIPool answer. The half-pixel difference was rounded away before any pooling happened. Every box whose corner falls anywhere inside a grid cell produces the same crop.

At output stride 16, one feature-grid cell is 16 image pixels. So RoIPool cannot distinguish boxes that differ by up to 16 pixels in the original photo. A box head can absorb that; a 28x28 mask cannot.

RoIAlign moved by exactly the amount the box moved. 2.5 → 3.0 and 4.5 → 5.0. Because the feature values encode their own column index, those numbers are directly readable as sampling positions. The half-pixel shift produced a half-unit shift.

sampling_ratio=2 means four sample points per output cell, bilinearly interpolated and averaged. Setting sampling_ratio=-1 makes the count adaptive to box size, which is the torchvision default inside Mask R-CNN.

The paste-back block shows the resolution cost. A 28x28 mask stretched to a 40x36 box means every mask pixel covers more than one image pixel. Boundaries are quantised to that grid, which is why Mask R-CNN masks look slightly blocky on large objects. PointRend and Mask R-CNN's later variants exist to fix precisely this.

The full model, without downloading weights#

import torch
from torchvision.models.detection import maskrcnn_resnet50_fpn_v2

torch.manual_seed(0)
model = maskrcnn_resnet50_fpn_v2(weights=None, weights_backbone=None, num_classes=3,
                                 min_size=128, max_size=128)
print("parameters:", f"{sum(p.numel() for p in model.parameters()):,}")

img = torch.rand(3, 128, 128)
mask = torch.zeros(1, 128, 128); mask[0, 30:80, 40:90] = 1
target = {"boxes": torch.tensor([[40., 30., 90., 80.]]),
          "labels": torch.tensor([1]),
          "masks": mask.to(torch.uint8)}

model.train()
losses = model([img], [target])
print("\nthe five losses Mask R-CNN trains on:")
for k, v in losses.items():
    print(f"   {k:20s} {v.item():.4f}")

model.eval()
with torch.no_grad():
    out = model([img])[0]
print("\nprediction dictionary from one image:")
for k, v in out.items():
    print(f"   {k:8s} {tuple(v.shape)}  dtype {v.dtype}")

parameters: 45,885,793

the five losses Mask R-CNN trains on:
   loss_classifier      1.0276
   loss_box_reg         0.0673
   loss_mask            1.5885
   loss_objectness      0.6989
   loss_rpn_box_reg     0.0075

prediction dictionary from one image:
   boxes    (100, 4)  dtype torch.float32
   labels   (100,)  dtype torch.int64
   scores   (100,)  dtype torch.float32
   masks    (100, 1, 128, 128)  dtype torch.float32

These weights are randomly initialised, so the boxes and scores carry no meaning. The loss values come from that random initialisation with manual_seed(0), and they will shift with a different torch version. What is stable and worth reading is the structure.

Five losses, trained together. Two belong to the region proposal network: loss_objectness (is there anything here?) and loss_rpn_box_reg (where exactly?). Three belong to the heads: class, box refinement and mask. A single backward pass updates all of them.

masks comes back at full image size, already pasted. torchvision does the 28x28-to-box stretch and the paste for you. Values are soft probabilities, so you threshold at 0.5 yourself.

One hundred detections from an untrained model is box_detections_per_img=100, the default cap, combined with a low score threshold. On a trained model you filter by scores > 0.5 or higher before doing anything else.

Targets need boxes, labels and masks together. masks must be uint8 with shape (N, H, W), one per instance, in the same order as boxes. Mismatched order is a silent, brutal bug.

Common mistakes#

Class 0 used for a real class. In torchvision detection models, 0 is reserved for background. num_classes=3 means background plus two real classes. Label your data from 1.

Passing an empty target. An image with no objects needs boxes shaped (0, 4) and masks shaped (0, H, W) — not None. Getting the dtype and shape right on empty tensors is a common training crash.

Feeding normalised tensors. These models normalise internally using their own transform. Pass images as float tensors in the range 0 to 1 and let the model do the rest.

Trusting the mask outside its box. The mask head cannot draw beyond its box, so a box that clips an arm produces a mask that clips the arm. When masks look truncated, inspect the boxes first.

Comparing mask AP against mIoU. They measure different things. See semantic, instance and panoptic segmentation for why the numbers do not move together.

Try it yourself#

Change output_size=2 to output_size=4 in the RoIAlign script and re-run. The resolution of the crop rises, and the half-pixel shift becomes visible in more cells. Then set sampling_ratio=1 and watch the interpolated values change, since one sample per cell is a coarser estimate of the same region.

What to learn next#

• Panoptic segmentation — merging these overlapping masks into one clean map.
• Object detection — the box stage this whole design rests on.
• YOLO — the single-stage tradition, and where its segmentation variants come from.