Mask R-CNN

Mask R-CNN is a model for instance segmentation introduced by Kaiming He, Georgia Gkioxari, Piotr Dollár and Ross Girshick at ICCV 2017 (arXiv, March 2017; journal version in IEEE TPAMI, 2020). It is a two-stage convolutional detector: Faster R-CNN plus a small fully convolutional branch that predicts a binary mask for each region, and RoIAlign, a feature-pooling layer without coordinate rounding. For each detected object it returns a box, a class, a score and a mask; with a different target, the same design detects human keypoints.

Problem Addressed

Faster R-CNN and fully convolutional networks (FCN; Long, Shelhamer & Darrell, 2015) were simple, strong baselines for detection and semantic segmentation; the authors wanted a comparable one for instance segmentation. The COCO winners of 2015 and 2016, the multi-stage cascade MNC (Dai, He & Sun, 2016) and FCIS (Li et al., 2017), were more complex, and FCIS made systematic errors on overlapping instances, by He et al.’s account; the instance segmentation article describes both.

Two things stood in the way of simply adding a mask output. RoI pooling rounds region coordinates to the feature grid, which classification tolerates but pixel-accurate masks do not; and FCN-style segmentation makes classes compete for each pixel through a softmax, which worked poorly for instances in the authors’ experiments.

Architecture

A region proposal network (RPN), unchanged from Faster R-CNN, proposes boxes from the backbone’s features; the second stage extracts a fixed-size feature for each box with RoIAlign and runs two heads on it in parallel, the box head (classification and box regression) and the mask head.

Component	ResNet-C4 configuration	ResNet-FPN configuration
Backbone	ResNet stages 1–4, stride 16	ResNet with a feature pyramid, levels $P_2$–$P_5$ ($P_6$ added for the RPN)
RoI features	RoIAlign, $7 \times 7 \times 1024$	RoIAlign from the level chosen by box size: $7 \times 7 \times 256$ (box), $14 \times 14 \times 256$ (mask)
Box head	ResNet stage 5 on each RoI, average pooling, class and box outputs	Two 1,024-d fully connected layers, class and box outputs
Mask head	From the stage-5 output: deconvolution to $14 \times 14 \times 256$, then $14 \times 14 \times K$	Four $3 \times 3$ convolutions (256 channels), deconvolution to $28 \times 28 \times 256$, then $1 \times 1$ convolution to $28 \times 28 \times K$

The table follows Figure 3 of the paper: deconvolutions are $2 \times 2$ with stride 2 (a transposed convolution), hidden layers use ReLU, and $K = 80$ on COCO.

Backbone

C4 takes features from the last layer of ResNet’s fourth stage, as in the original ResNet detector (He et al., 2016); FPN uses a feature pyramid network, which assigns each RoI to a level by its size. The networks were ResNet-50, ResNet-101 and ResNeXt-101 (64×4d; Xie et al., 2017). The authors advise against C4 in practice: its per-RoI stage 5 is heavy, and FPN was both more accurate and faster.

RoIAlign

RoIAlign keeps region and bin boundaries at their real-valued positions ($x/16$ rather than a rounded $[x/16]$), interpolates features bilinearly at four regular points in each bin, and averages them (max works as well). The two-stage object detectors article derives it and compares it with RoI pooling.

Mask Head

For each RoI the mask head outputs $K m^2$ values, one $m \times m$ mask per class, with $m = 28$ in the FPN head and $14$ in the C4 head. Being convolutional, it keeps the RoI’s spatial layout in every layer instead of collapsing it into a vector, as the fully connected mask predictors of DeepMask and MNC did.

Training

The loss on each sampled RoI is

$$L = L_{\text{cls}} + L_{\text{box}} + L_{\text{mask}},$$

where $L_{\text{cls}}$ and $L_{\text{box}}$ are Fast R-CNN’s classification and box regression losses (Girshick, 2015). For a positive RoI of ground-truth class $k$, with mask logits $z_{k,ij}$ and binary target $t_{ij}$ on the $m \times m$ grid,

$$L_{\text{mask}} = -\frac{1}{m^2} \sum_{i,j} \Big[ t_{ij} \log \sigma(z_{k,ij}) + (1 - t_{ij}) \log\big(1 - \sigma(z_{k,ij})\big) \Big],$$

a per-pixel sigmoid with average binary cross-entropy. Only the $k$-th mask contributes; the other $K - 1$ outputs receive no gradient from this RoI. The authors call this decoupling mask and class prediction: masks of different classes do not compete, and the box head picks the class. On COCO minival with ResNet-50-C4, replacing it with a per-pixel softmax and multinomial loss, as in FCN, dropped mask AP from 30.3 to 24.8. Predicting a single class-agnostic mask gave 29.7.

Targets and sampling. An RoI is positive if its IoU with a ground-truth box is at least 0.5, and $L_{\text{mask}}$ is defined only on positives. The mask target is the intersection of the RoI with the ground-truth mask of its matched object, resampled to $m \times m$. Training is image-centric: images are resized to a shorter side of 800 pixels, each mini-batch holds 2 images per GPU, and each image contributes $N$ sampled RoIs at a 1:3 ratio of positives to negatives, with $N = 64$ for C4 and 512 for FPN.

Schedule. On 8 GPUs (16 images per mini-batch), SGD ran for 160k iterations at learning rate 0.02, divided by 10 at 120k, with weight decay 0.0001 and momentum 0.9. The RPN used anchors of 5 scales and 3 aspect ratios. For convenient ablation, the RPN was trained separately, without sharing features, unless stated otherwise. Training on COCO trainval35k took 32 hours with ResNet-50-FPN and 44 with ResNet-101-FPN on one 8-GPU machine. The arXiv v3 appendix adds that end-to-end training, ignoring gradients with respect to RoI coordinates as in Faster R-CNN’s approximate joint training, gained 0.6 mask AP and 1.2 box AP (ResNeXt-101-FPN, minival).

Example: mask targets and box errors

The NumPy and SciPy code draws a cup (body and handle) whose box also holds part of a touching cup, builds the $28 \times 28$ target by RoIAlign-style sampling of the object’s own mask in a proposal box, and binarizes it. A smoothed copy of the target stands in for a predicted mask: it is pasted into the box at threshold 0.5 and compared with the ground truth as the proposal shifts.

import numpy as np
from scipy import ndimage

# A 160 x 160 image with two touching objects of the same class ("cups"):
# the instance we train on (body + handle) and a neighbor that overlaps its box.
H = W = 160
yy, xx = np.mgrid[:H, :W] + 0.5                       # pixel centers
body = (np.abs(yy - 80) <= 30) & (np.abs(xx - 70) <= 22)
handle = ((yy - 80) ** 2 + (xx - 96) ** 2 <= 14 ** 2) & ((yy - 80) ** 2 + (xx - 96) ** 2 >= 8 ** 2) & (xx > 92)
gt = body | handle                                    # ground-truth mask of our instance
neighbor = ((yy - 60) ** 2 + (xx - 122) ** 2 <= 20 ** 2) & ~gt
img = 0.2 + 0.6 * gt + 0.6 * neighbor + np.random.default_rng(0).normal(0, 0.05, (H, W))

def box_of(m):                                        # tight box (x0, y0, x1, y1), pixel edges
    r, c = np.nonzero(m)
    return np.array([c.min(), r.min(), c.max() + 1, r.max() + 1], float)

def roi_align(f, box, m=28, s=2):
    """Average of s x s bilinear samples in each of m x m bins; no rounding of the box."""
    x0, y0, x1, y1 = box
    t = (np.arange(m * s) + 0.5) / (m * s)
    ys, xs = np.meshgrid(y0 + t * (y1 - y0) - 0.5, x0 + t * (x1 - x0) - 0.5, indexing="ij")
    v = ndimage.map_coordinates(f.astype(float), [ys, xs], order=1, mode="constant")
    return v.reshape(m, s, m, s).mean(axis=(1, 3))

def paste(mm, box, thr=0.5):
    """Resize an m x m soft mask to the box (bilinear), threshold, place it in the image."""
    x0, y0, x1, y1 = box
    m = mm.shape[0]
    out = np.zeros((H, W), bool)
    r0, r1 = int(np.floor(y0)), int(np.ceil(y1))
    c0, c1 = int(np.floor(x0)), int(np.ceil(x1))
    py, px = np.mgrid[r0:r1, c0:c1] + 0.5
    gy, gx = (py - y0) / (y1 - y0) * m - 0.5, (px - x0) / (x1 - x0) * m - 0.5
    inside = (py > y0) & (py < y1) & (px > x0) & (px < x1)
    v = ndimage.map_coordinates(mm, [gy, gx], order=1, mode="nearest")
    out[r0:r1, c0:c1] = (v >= thr) & inside
    return out

def iou(a, b):
    return (a & b).sum() / (a | b).sum()

g = box_of(gt)
print("GT box (x0, y0, x1, y1):", g, " GT pixels:", gt.sum())
print("neighbor pixels inside the GT box:", (neighbor & (xx > g[0]) & (xx < g[2]) & (yy > g[1]) & (yy < g[3])).sum())

# Mask target: the instance's own mask (not the neighbor's) sampled in the proposal.
T_soft = roi_align(gt, g)
T = T_soft >= 0.5
print("28 x 28 target from the exact box: %d of 784 cells foreground" % T.sum())
print("image patch in the same box, bright cells (> 0.5): %d" % (roi_align(img, g) > 0.5).sum())

# A stand-in "prediction": the soft target, smoothed as a network output would be.
print("\nshift (px)  GT inside box  target cells  mask IoU after paste")
for d in (0, 4, 8, 16):
    box = g + np.array([d, d, d, d])                  # proposal shifted right and down
    inside = (gt & (xx > box[0]) & (xx < box[2]) & (yy > box[1]) & (yy < box[3])).sum() / gt.sum()
    T = roi_align(gt, box) >= 0.5                     # intersection of RoI and GT mask
    pred = ndimage.gaussian_filter(T.astype(float), 0.7)
    print(f"{d:6d}      {inside:12.3f}  {T.sum():12d}  {iou(paste(pred, box), gt):10.3f}")

# The same shifts applied to a box that is 20% larger than the object on each side.
print("\nproposal enlarged by 20% per side")
for d in (0, 8, 16):
    w, h = g[2] - g[0], g[3] - g[1]
    box = g + np.array([-0.2 * w, -0.2 * h, 0.2 * w, 0.2 * h]) + d
    T = roi_align(gt, box) >= 0.5
    pred = ndimage.gaussian_filter(T.astype(float), 0.7)
    print(f"shift {d:2d}: target cells {T.sum():3d}, mask IoU {iou(paste(pred, box), gt):.3f}")

Output:

GT box (x0, y0, x1, y1): [ 48.  50. 110. 110.]  GT pixels: 2894
neighbor pixels inside the GT box: 156
28 x 28 target from the exact box: 614 of 784 cells foreground
image patch in the same box, bright cells (> 0.5): 644

shift (px)  GT inside box  target cells  mask IoU after paste
     0             1.000           614       0.987
     4             0.862           521       0.850
     8             0.735           438       0.703
    16             0.513           322       0.496

proposal enlarged by 20% per side
shift  0: target cells 310, mask IoU 0.970
shift  8: target cells 325, mask IoU 0.916
shift 16: target cells 274, mask IoU 0.838

The patch has 644 bright cells but the target only 614 foreground cells: the difference is the neighbor, background here because the target comes from this instance’s mask alone, so the head must tell its own object from a same-class neighbor by shape and position (over noise seeds 0–9 the bright-cell count ranged from 644 to 648; nothing else changed). A pasted mask cannot extend past its box, so a 16-pixel shift of a 62-pixel box loses half the object and mask IoU falls to 0.496. A box 20% larger on each side keeps more of the object in view, at the cost of a coarser grid (310 foreground cells). In our reading, mask quality is bounded by box quality, one reason to run the mask head on refined detection boxes at test time.

Example: why per-class sigmoids

Here one positive RoI of class “bowl” has mask logits for “cup” and “bowl” that both outline the object, as for two similar-looking classes. The code compares Mask R-CNN’s loss with an FCN-style per-pixel softmax over background plus the three classes (the paper does not say whether its softmax variant had a background channel), and checks both gradients against finite differences.

import numpy as np

# One positive RoI of ground-truth class k = 1 ("bowl"), with a 14 x 14 mask target:
# a disk, as RoIAlign would sample it from the instance's ground-truth mask.
m, K, k = 14, 3, 1
yy, xx = np.mgrid[:m, :m] + 0.5
T = ((yy - 7) ** 2 + (xx - 7) ** 2 <= 5.5 ** 2).astype(float)

# Mask logits for K = 3 classes ("cup", "bowl", "vase"). Cups and bowls look alike,
# so channels 0 and 1 both predict the object's shape; channel 2 predicts nothing.
shape = np.where(T > 0, 4.0, -4.0)
Z = np.stack([shape, shape, np.full((m, m), -4.0)])
Zbg = -shape                                          # background logit for the softmax variant

sig = lambda z: 1 / (1 + np.exp(-z))

def loss_sigmoid(Z):
    """Mask R-CNN: average binary cross-entropy on the ground-truth class's channel only."""
    p = sig(Z[k])
    return -np.mean(T * np.log(p) + (1 - T) * np.log(1 - p))

def grad_sigmoid(Z):
    G = np.zeros_like(Z)
    G[k] = (sig(Z[k]) - T) / T.size                   # other channels get no gradient
    return G

def loss_softmax(Z, Zbg):
    """FCN-style: per-pixel softmax over background + K classes, multinomial cross-entropy."""
    A = np.concatenate([Zbg[None], Z])                # channel 0 = background
    A = A - A.max(axis=0)
    P = np.exp(A) / np.exp(A).sum(axis=0)
    label = np.where(T > 0, k + 1, 0)                 # class index per pixel
    return -np.mean(np.log(np.take_along_axis(P, label[None], 0))), P

def grad_softmax(Z, Zbg):
    _, P = loss_softmax(Z, Zbg)
    Y = np.zeros_like(P)
    np.put_along_axis(Y, np.where(T > 0, k + 1, 0)[None], 1.0, 0)
    return (P - Y)[1:] / T.size                       # gradient for the K class channels

# Check both hand-written gradients against central finite differences.
rng = np.random.default_rng(0)
Zr, Zb = rng.normal(0, 2, Z.shape), rng.normal(0, 2, (m, m))
fd_s, fd_m = np.zeros_like(Zr), np.zeros_like(Zr)
for idx in np.ndindex(Zr.shape):
    E = np.zeros_like(Zr); E[idx] = 1e-6
    fd_s[idx] = (loss_sigmoid(Zr + E) - loss_sigmoid(Zr - E)) / 2e-6
    fd_m[idx] = (loss_softmax(Zr + E, Zb)[0] - loss_softmax(Zr - E, Zb)[0]) / 2e-6
print("max |analytic - finite difference|: sigmoid %.1e, softmax %.1e"
      % (abs(grad_sigmoid(Zr) - fd_s).max(), abs(grad_softmax(Zr, Zb) - fd_m).max()))

Ls, (Lm, P) = loss_sigmoid(Z), loss_softmax(Z, Zbg)
Gs, Gm = grad_sigmoid(Z), grad_softmax(Z, Zbg)
print("\n                     sigmoid (Mask R-CNN)   softmax (FCN-style)")
print("loss                 %10.4f             %10.4f" % (Ls, Lm))
for c, name in enumerate(["cup", "bowl", "vase"]):
    print("|grad| channel %-5s %10.1e             %10.1e" % (name, np.linalg.norm(Gs[c]), np.linalg.norm(Gm[c])))
obj = T > 0
print("on object pixels: p(cup) + p(bowl) = %.3f (sigmoid), %.3f (softmax)"
      % ((sig(Z[0]) + sig(Z[1]))[obj].mean(), (P[1] + P[2])[obj].mean()))
print("p(bowl) on object pixels: %.3f (sigmoid), %.3f (softmax)" % (sig(Z[1])[obj].mean(), P[2][obj].mean()))

# Inference selects the mask of the class the box head predicts.
for c, name in enumerate(["cup", "bowl"]):
    ms = sig(Z[c]) >= 0.5
    print(f"box head says {name:4s}: sigmoid mask IoU {(ms & obj).sum() / (ms | obj).sum():.3f}")

Output:

max |analytic - finite difference|: sigmoid 1.6e-10, softmax 4.1e-10

                     sigmoid (Mask R-CNN)   softmax (FCN-style)
loss                     0.0181                 0.3119
|grad| channel cup      0.0e+00                2.4e-02
|grad| channel bowl     1.3e-03                2.4e-02
|grad| channel vase     0.0e+00                1.9e-05
on object pixels: p(cup) + p(bowl) = 1.964 (sigmoid), 1.000 (softmax)
p(bowl) on object pixels: 0.982 (sigmoid), 0.500 (softmax)
box head says cup : sigmoid mask IoU 1.000
box head says bowl: sigmoid mask IoU 1.000

Under the sigmoid loss the cup channel’s shape costs nothing and gets zero gradient. Under the softmax, cup and bowl probabilities on the object must sum to at most 1, so the two plausible masks split it at 0.5 each, the loss is 17 times larger, and the gradient pushes the cup channel down: the mask head is made to classify, which the box head already does. This illustrates the decoupling argument; the 5.5-point COCO gap above is the authors’ evidence.

Inference

At test time the RPN supplies 300 proposals (C4) or 1,000 (FPN). The box branch scores and refines them, non-maximum suppression removes duplicates (threshold 0.3 by default, according to the arXiv v3 appendix), and the mask branch runs only on the 100 highest-scoring detections. Unlike training, this runs the heads in sequence, which the authors found faster and more accurate, as masks are computed on fewer, better-localized boxes. Only the mask of the predicted class $k$ is used: the $m \times m$ probabilities are resized to the box and binarized at 0.5. The mask’s score is the box’s classification confidence, as Huang et al. (2019) note in questioning the practice.

The mask branch adds about 20% to Faster R-CNN’s cost. A ResNet-101-FPN model with shared RPN features ran at 195 ms per image on an Nvidia Tesla M40 GPU, plus 15 ms of CPU time to resize outputs (about 5 frames per second); the C4 variant took about 400 ms.

Performance

COCO test-dev results from He et al. (2017), all single models (Mask R-CNN trained on trainval35k); ablations use the 5k-image minival split.

Model	Backbone	Split	Mask AP	Box AP
MNC (2015 winner)	ResNet-101-C4	test-dev	24.6	–
FCIS + OHEM (2016 winner)	ResNet-101-C5-dilated	test-dev	29.2	–
FCIS+++ (multi-scale, flip, OHEM)	ResNet-101-C5-dilated	test-dev	33.6	–
Mask R-CNN	ResNet-101-C4	test-dev	33.1	–
Mask R-CNN	ResNet-101-FPN	test-dev	35.7	38.2
Mask R-CNN	ResNeXt-101-FPN	test-dev	37.1	39.8

The ResNet-101-FPN model’s 35.7 mask AP (AP@0.5 58.0, AP@0.75 37.8) comes with 15.5 AP on small objects and 52.4 on large ones. Trained without the mask branch, the same system reached 37.3 box AP, 0.9 below Mask R-CNN; the authors attribute that gap to multi-task training.

Minival ablations, mask AP:

• Backbone: 30.3 (ResNet-50-C4), 32.7 (ResNet-101-C4), 33.6 (ResNet-50-FPN), 35.4 (ResNet-101-FPN), 36.7 (ResNeXt-101-FPN).
• RoI layer (ResNet-50-C4): RoIAlign 30.3, RoI pooling 26.9, RoIWarp (bilinear but rounded) 27.2; with stride-32 C5 features, +7.3 (30.9 against 23.6); with FPN, +1.5.
• Mask head (ResNet-50-FPN): a fully convolutional head gave 33.6, against 31.5 for a fully connected one.

On Cityscapes (arXiv v3 appendix, ResNet-50-FPN, fine annotations only), Mask R-CNN reported 26.2 AP on the test set and 31.5 on val, and 32.0 test AP when initialized from a COCO model. The previous best entry, DIN, had 20.0 and the concurrent SGN 25.0, both trained on fine and coarse data. A second appendix raises COCO minival mask AP from 36.7 to 41.8 with end-to-end training, ImageNet-5k pre-training, ResNeXt-152, non-local blocks, longer schedules and train- and test-time augmentation.

Strengths

• Simple extension. One head and an aligned pooling layer turn a detector into an instance segmenter at about 20% extra cost.
• Decoupled tasks. The mask head only separates figure from ground; class-agnostic masks lose 0.6 AP.
• Multi-task gain. Mask supervision raised box AP by 0.9.
• Generality. Keypoints need only a new target and loss.

Limitations

• Mask resolution. A $28 \times 28$ grid stretched over a large box blurs boundaries (measured in the instance segmentation article); PointRend’s authors note that COCO’s often coarse ground truth may not fully reflect gains in mask quality.
• Masks confined to boxes. Localization errors become mask errors, as the example shows.
• Score is not mask quality. The classification score ranks masks whatever their IoU.
• Speed. About 5 frames per second on 2017 hardware, with per-RoI computation and hand-set NMS.
• Overlapping outputs. A single label map, as panoptic segmentation requires, needs an extra merging step.

Historical Importance

Mask R-CNN received the Best Paper Award (Marr Prize) at ICCV 2017, according to the PAMI Technical Committee’s list of ICCV paper awards. It became a common baseline: Mask Scoring R-CNN and PointRend, for example, build on it and measure against it. According to the arXiv v3 appendix, the first three teams of the COCO 2017 instance segmentation challenge reportedly all extended it.

Variants

• Human keypoints (He et al., 2017). Each of the $K$ keypoint types becomes a one-hot $m \times m$ mask with a single foreground pixel, trained with an $m^2$-way softmax cross-entropy per visible keypoint. The head (eight $3 \times 3$ convolutions of 512 channels, a deconvolution and 2× bilinear upsampling) outputs $56 \times 56$, as keypoints need finer localization. With ResNet-50-FPN it reached 62.7 keypoint AP on COCO test-dev, and 63.1 when trained jointly with person masks, against 61.8 for the 2016 challenge winner. RoIAlign added 4.4 points over RoI pooling (minival). TorchVision calls it keypointrcnn_resnet50_fpn.
• Mask Scoring R-CNN (Huang et al., 2019). A MaskIoU head takes the RoI features and the predicted mask and regresses the mask’s IoU with its ground truth; the mask score is the classification score times this estimate. On COCO test-dev with ResNet-101-FPN, mask AP rose from 37.0 for the authors’ Mask R-CNN baseline to 38.3.
• PointRend (Kirillov et al., 2020). A coarse $7 \times 7$ mask is refined to $224 \times 224$ in five subdivision steps, each re-predicting up to 784 of the most uncertain points with a small MLP on fine $P_2$ features and the coarse prediction. With ResNet-50-FPN it raised COCO val2017 mask AP from 35.2 to 36.3 (37.6 to 39.7 against the finer LVIS annotations of COCO’s categories) and Cityscapes val AP from 33.0 to 35.8.
• Cascade Mask R-CNN. Cai and Vasconcelos (2021) add a mask branch to one or every stage of the multi-stage Cascade R-CNN; the instance segmentation article covers it and Hybrid Task Cascade.

Implementations

The authors released Mask R-CNN in Detectron (Girshick et al., 2018), a Caffe2 system now deprecated; its PyTorch successor Detectron2 (Wu et al., 2019) has Mask R-CNN baselines with C4, dilated-C5 and FPN backbones, and PointRend as a project. TorchVision’s maskrcnn_resnet50_fpn is listed in its documentation (version 0.29) at 34.6 mask AP on COCO val2017, and maskrcnn_resnet50_fpn_v2, with an improved recipe, at 41.8; both return soft image-sized masks to threshold at 0.5. The API changes between releases, so check the installed version’s documentation.

The torch-only snippet below, run with PyTorch 2.14.1, builds the FPN mask head from Figure 3 and selects each RoI’s ground-truth channel for the loss.

import torch
from torch import nn

# The FPN mask head of Mask R-CNN (He et al., 2017, Figure 3, right), for K = 80 COCO classes.
torch.manual_seed(0)
K = 80
layers = []
for _ in range(4):                                    # four 3x3 convs, 256 channels, ReLU
    layers += [nn.Conv2d(256, 256, 3, padding=1), nn.ReLU()]
layers += [nn.ConvTranspose2d(256, 256, 2, stride=2), nn.ReLU(),   # 2x2 deconv, stride 2
           nn.Conv2d(256, K, 1)]                      # 1x1 conv: one mask logit map per class
head = nn.Sequential(*layers)

x = torch.randn(100, 256, 14, 14)                     # RoIAlign features for 100 detections
for layer in head:
    x = layer(x)
    if not isinstance(layer, nn.ReLU):
        n = sum(p.numel() for p in layer.parameters())
        print(f"{layer.__class__.__name__:16s} -> {tuple(x.shape)}  params {n:,}")
print(f"total parameters: {sum(p.numel() for p in head.parameters()):,}")

# Training uses only the ground-truth class's channel of each positive RoI.
labels = torch.randint(0, K, (100,))
targets = (torch.rand(100, 28, 28) > 0.5).float()     # stand-in for 28 x 28 mask targets
logits = x[torch.arange(100), labels]                 # (100, 28, 28)
loss = nn.functional.binary_cross_entropy_with_logits(logits, targets)
loss.backward()
print("selected logits:", tuple(logits.shape), " loss: %.3f" % loss.item())

# For comparison: the last two layers of the fully connected mask head in the ablation.
mlp_tail = 1024 * 1024 + 1024 + 1024 * K * 28 * 28 + K * 28 * 28
print(f"fc 1024 -> 1024 -> 80*28*28 (last two layers only): {mlp_tail:,} parameters")
print("torch", torch.__version__)

Output:

Conv2d           -> (100, 256, 14, 14)  params 590,080
Conv2d           -> (100, 256, 14, 14)  params 590,080
Conv2d           -> (100, 256, 14, 14)  params 590,080
Conv2d           -> (100, 256, 14, 14)  params 590,080
ConvTranspose2d  -> (100, 256, 28, 28)  params 262,400
Conv2d           -> (100, 80, 28, 28)  params 20,560
total parameters: 2,643,280
selected logits: (100, 28, 28)  loss: 0.693
fc 1024 -> 1024 -> 80*28*28 (last two layers only): 65,337,600 parameters
torch 2.14.1

The convolutional head has 2.6 million parameters; by our count, the last two layers of the ablation’s fully connected head have 65 million, consistent with the paper’s remark that the convolutional head needs fewer. The initial loss is $\ln 2 \approx 0.693$ because untrained logits are near zero.