Prerequisites
This notebook reuses tables created by other example notebooks. Run them first:
Fine-tune Mask2Former on Pascal VOC and collect metrics¶
Fine-tune a Mask2Former model for semantic segmentation on the Pascal VOC 2012 tables, and collect per-sample predictions and IoU back into a 3LC run.
We fine-tune Mask2Former on the tables from create-pascal-voc-semseg-table. Mask2Former is a universal segmentation model — it predicts a set of binary masks plus a class for each. We use it in semantic mode: starting from an ADE20k-semantic checkpoint (its Swin backbone and decoders already speak segmentation), we swap in a 21-class head for VOC and fine-tune, then collapse the mask set back to a dense per-pixel
label map with the processor’s post_process_semantic_segmentation. Only the classification head is reinitialized, so we fine-tune it at a higher learning rate than the pretrained weights. The void boundary (id 255) is excluded from both the loss and the metrics. We track train/val loss every epoch and, after training, collect per sample on both splits: the predicted segmentation (as RLE), mean IoU, the per-image confusion matrix (via the core metrics helper), and a pooled decoder embedding
(reduced to 2D with PaCMAP).
Project setup¶
[ ]:
# Parameters
PROJECT_NAME = "3LC Tutorials - Pascal VOC 2012"
DATASET_NAME = "pascal-voc-2012"
RUN_NAME = "Fine-tune Mask2Former"
RUN_DESCRIPTION = "Mask2Former (semantic mode) fine-tuned on Pascal VOC 2012"
CHECKPOINT = "facebook/mask2former-swin-tiny-ade-semantic"
EPOCHS = 10
BATCH_SIZE = 4
LR = 5e-5 # learning rate for the pretrained weights (backbone + decoders)
HEAD_LR_MULT = 10 # the reinitialized classification head trains at LR * HEAD_LR_MULT
IMAGE_SIZE = 384
NUM_WORKERS = 0 # 0 is safest in notebooks: notebook-defined transforms can't be pickled to
# spawned DataLoader workers. Bump it only when running as an importable module.
SEED = 42
INSTALL_DEPENDENCIES = True
Install dependencies¶
[ ]:
if INSTALL_DEPENDENCIES:
%pip install -q 3lc transformers torch torchvision tqdm pacmap
Imports¶
[ ]:
import random
import numpy as np
import tlc
import torch
from tlc.data_types import SemanticSegmentation
from tlc.integration.torch.samplers import create_sampler
from tlc.metrics.semantic_segmentation import semantic_segmentation_metrics
from tlc.schemas import SemanticSegmentationRleSchema
from torch.utils.data import DataLoader
from tqdm.auto import tqdm
from transformers import AutoImageProcessor, Mask2FormerForUniversalSegmentation
Class universe¶
VOC has 21 classes (background + 20 objects); void (id 255) is the ignore boundary. background (id 0) is not a value-map class — it rides in the column schema’s metadata — so both the GT and predicted value maps show the 20 objects, with the GT map additionally carrying the tagged void class. id2label / label2id are what we hand to the model’s reinitialized classification head (the model still predicts background as id 0).
[ ]:
VOC_CLASS_NAMES = [
"background",
"aeroplane",
"bicycle",
"bird",
"boat",
"bottle",
"bus",
"car",
"cat",
"chair",
"cow",
"diningtable",
"dog",
"horse",
"motorbike",
"person",
"pottedplant",
"sheep",
"sofa",
"train",
"tvmonitor",
]
BACKGROUND_ID = 0
VOID_ID = 255
NUM_CLASSES = len(VOC_CLASS_NAMES) # 21
GT_CLASSES = {**{i: n for i, n in enumerate(VOC_CLASS_NAMES)}, VOID_ID: "void"}
PREDICTED_CLASSES = {i: n for i, n in enumerate(VOC_CLASS_NAMES)}
id2label = {i: n for i, n in enumerate(VOC_CLASS_NAMES)}
label2id = {n: i for i, n in enumerate(VOC_CLASS_NAMES)}
Load the processor and model¶
The processor turns images (and, for training, segmentation maps) into the mask-classification targets Mask2Former expects, and post-processes predictions back into dense label maps. We set ignore_index to the VOC void id and keep do_reduce_labels=False (VOC’s background is a real class, not an ignore label).
The model loads from an ADE20k-semantic checkpoint; num_labels=21 with ignore_mismatched_sizes=True reinitializes the classification head for VOC while keeping the pretrained backbone and pixel decoder.
[ ]:
processor = AutoImageProcessor.from_pretrained(
CHECKPOINT,
size={"height": IMAGE_SIZE, "width": IMAGE_SIZE},
ignore_index=VOID_ID,
do_reduce_labels=False,
)
model = Mask2FormerForUniversalSegmentation.from_pretrained(
CHECKPOINT,
num_labels=NUM_CLASSES,
id2label=id2label,
label2id=label2id,
ignore_mismatched_sizes=True,
)
Sample transform¶
We feed the table to the DataLoader via Table.with_transform rather than a custom Dataset. The transform returns a (PIL image, dense label map) pair per row; the Mask2Former processor handles resizing and target construction in the collate function. Curation is handled by the sampler (create_sampler drops zero-weight rows), so Dashboard exclusions take effect on the next run.
[ ]:
def voc_sample_transform(sample: dict) -> tuple:
seg: SemanticSegmentation = sample["mask"]
return sample["image"].convert("RGB"), seg.mask.astype(np.int64)
def collate_fn(batch: list) -> dict:
images = [image for image, _ in batch]
masks = [mask for _, mask in batch]
return processor(images=images, segmentation_maps=masks, return_tensors="pt")
Metrics collection¶
For each sample we predict at the original image size, then derive the dense label map with post_process_semantic_segmentation. We store the predicted segmentation (as RLE), the mean IoU and the per-image confusion matrix (both from the core helper, void excluded), and a pooled embedding from the transformer decoder’s per-query hidden states.
[ ]:
def collect_metrics(run: tlc.Run, table: tlc.Table, model, device: torch.device, epoch: int) -> float:
predictions: list[np.ndarray] = []
ious: list[float] = []
confusion_matrices: list[list[int]] = []
embeddings: list[np.ndarray] = []
model.eval()
with torch.no_grad():
for idx in tqdm(range(len(table)), desc="collect", leave=False):
row = table[idx]
image = row["image"].convert("RGB")
seg: SemanticSegmentation = row["mask"]
width, height = image.size
inputs = processor(images=image, return_tensors="pt").to(device)
outputs = model(**inputs)
pred_map = (
processor.post_process_semantic_segmentation(outputs, target_sizes=[(height, width)])[0]
.cpu()
.numpy()
.astype(np.int32)
)
# Pooled decoder embedding: mean over the object queries -> one vector per image.
decoder_hidden = outputs.transformer_decoder_last_hidden_state # (1, num_queries, hidden)
embeddings.append(decoder_hidden.mean(dim=1).squeeze(0).cpu().numpy().astype(np.float32))
# Bare label map; the metrics writer serializes it via the predicted_segmentation
# schema below (which records background id 0 in its metadata), so no wrapper is needed.
predictions.append(pred_map)
m = semantic_segmentation_metrics(
pred_map, seg.mask, GT_CLASSES, background=BACKGROUND_ID, void=VOID_ID, include_background=True
)
ious.append(m["mean_iou"])
confusion_matrices.append([int(x) for cm_row in m["confusion_matrix"] for x in cm_row])
run.add_metrics(
{
"predicted_segmentation": predictions,
"iou": ious,
"confusion_matrix": confusion_matrices,
"embedding": embeddings,
"epoch": [epoch] * len(ious),
},
schema={
# Background (id 0) rides in the schema metadata, not the value map.
"predicted_segmentation": SemanticSegmentationRleSchema(
classes=PREDICTED_CLASSES, background=BACKGROUND_ID
),
"embedding": tlc.schemas.EmbeddingSchema(shape=len(embeddings[0])),
},
foreign_table_url=table.url,
)
# Average only over images with a defined IoU; a degenerate image yields nan, and a nan headline
# would be logged into the run object as an invalid JSON token.
finite = [v for v in ious if np.isfinite(v)]
return float(np.mean(finite)) if finite else float("nan")
Reproducibility helpers¶
[ ]:
def seed_everything(seed: int) -> None:
random.seed(seed)
np.random.seed(seed)
torch.manual_seed(seed)
torch.cuda.manual_seed_all(seed)
Load the tables and initialize the Run¶
.latest() picks up the newest revision of each table, so retraining consumes any Dashboard curation without code changes.
[ ]:
seed_everything(SEED)
device = torch.device("cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu")
print(f"Using device: {device}")
model = model.to(device)
train_table = tlc.Table.from_names(table_name="train", dataset_name=DATASET_NAME, project_name=PROJECT_NAME).latest()
val_table = tlc.Table.from_names(table_name="val", dataset_name=DATASET_NAME, project_name=PROJECT_NAME).latest()
print(f"Using train table {train_table.url}")
print(f"Using val table {val_table.url}")
# with_transform returns a map-style view — pass it straight to a DataLoader. create_sampler reads the
# weight column: zero-weight rows (excluded in the Dashboard) are dropped, train is shuffled, val is
# sequential.
train_view = train_table.with_transform(voc_sample_transform)
val_view = val_table.with_transform(voc_sample_transform)
train_sampler = create_sampler(train_table, weighted=False, shuffle=True)
val_sampler = create_sampler(val_table, weighted=False, shuffle=False)
print(f"Train: {len(train_sampler)} of {len(train_table)} rows after weight filtering | Val: {len(val_sampler)}")
run = tlc.init(
PROJECT_NAME,
run_name=RUN_NAME,
description=RUN_DESCRIPTION,
parameters={
"checkpoint": CHECKPOINT,
"epochs": EPOCHS,
"batch_size": BATCH_SIZE,
"lr": LR,
"image_size": IMAGE_SIZE,
"seed": SEED,
},
)
train_loader = DataLoader(
train_view,
batch_size=BATCH_SIZE,
sampler=train_sampler,
num_workers=NUM_WORKERS,
collate_fn=collate_fn,
)
val_loader = DataLoader(
val_view,
batch_size=BATCH_SIZE,
sampler=val_sampler,
num_workers=NUM_WORKERS,
collate_fn=collate_fn,
)
Train¶
Two parameter groups give a discriminative learning rate: the pretrained backbone and decoders update gently at LR, while the freshly-initialized classification head trains at LR * HEAD_LR_MULT so it can catch up to weights that already learned good segmentation features on ADE20k. Each epoch logs train and val loss; the full per-sample metrics collection runs once after the loop.
[ ]:
# Discriminative LR: the reinitialized class head learns faster than the pretrained backbone/decoders.
head_params = list(model.class_predictor.parameters())
head_param_ids = {id(p) for p in head_params}
pretrained_params = [p for p in model.parameters() if id(p) not in head_param_ids]
optimizer = torch.optim.AdamW(
[
{"params": pretrained_params, "lr": LR}, # pretrained backbone + pixel/transformer decoders
{"params": head_params, "lr": LR * HEAD_LR_MULT}, # freshly initialized for VOC's 21 classes
]
)
for epoch in range(EPOCHS):
model.train()
train_loss = 0.0
n_batches = 0
for batch in tqdm(train_loader, desc=f"epoch {epoch}", leave=False):
optimizer.zero_grad()
outputs = model(
pixel_values=batch["pixel_values"].to(device),
mask_labels=[m.to(device) for m in batch["mask_labels"]],
class_labels=[c.to(device) for c in batch["class_labels"]],
)
loss = outputs.loss
loss.backward()
optimizer.step()
train_loss += loss.item()
n_batches += 1
train_loss /= max(n_batches, 1)
# Cheap val-loss pass every epoch (same loss the model trains on); the expensive per-sample metrics
# collection runs once after training.
model.eval()
val_loss = 0.0
n_val_batches = 0
with torch.no_grad():
for batch in val_loader:
outputs = model(
pixel_values=batch["pixel_values"].to(device),
mask_labels=[m.to(device) for m in batch["mask_labels"]],
class_labels=[c.to(device) for c in batch["class_labels"]],
)
val_loss += outputs.loss.item()
n_val_batches += 1
val_loss /= max(n_val_batches, 1)
log_entry = {"epoch": epoch, "train_loss": train_loss, "val_loss": val_loss}
tlc.log(log_entry)
print(" ".join(f"{k}={v:.4f}" if k != "epoch" else f"epoch {v}" for k, v in log_entry.items()))
Collect metrics and reduce embeddings¶
With training finished, we run the full per-sample collection once on both the val and train splits, then fit one PaCMAP model and apply it across every metrics table so they share a single 2D space.
[ ]:
print("Collecting per-sample metrics on val and train...")
collect_metrics(run, val_table, model, device, epoch=EPOCHS - 1)
collect_metrics(run, train_table, model, device, epoch=EPOCHS - 1)
print("Reducing embeddings (PaCMAP)...")
run.reduce_embeddings_by_foreign_table_url(val_table.url, method="pacmap", delete_source_tables=True)
run.set_status_completed()
print(f"Run: {run.url}")
Next steps¶
Open the Run in the 3LC Dashboard to overlay predictions on the images, sort by IoU to find the hardest classes and images, and use the per-image confusion matrix to see which VOC classes get confused. Exclude or relabel samples in the Dashboard, then re-run — .latest() picks up the curated revision automatically.