ADE20K probe on CanViT's 32 × 32 canvas

A linear ADE20K semantic segmentation probe on CanViT's 32 × 32 canvas, trained on frozen features with the paper's probing protocol.

CanViT, the Canvas Vision Transformer, is an active-vision foundation model: it sees a scene through a sequence of glimpses and remembers it on a scene-wide canvas.

Paper (NeurIPS 2026) · Code · Project page · All checkpoints

Usage

uv add "canvit-pytorch>=0.2"

import torch
from PIL import Image
from canvit_pytorch import CanViTForSemanticSegmentation, Viewpoint, sample_at_viewpoint
from canvit_pytorch.benchmarks.ade20k import CLASS_NAMES
from canvit_pytorch.preprocess import preprocess

model = CanViTForSemanticSegmentation.from_pretrained_with_probe(
    pretrained_repo="canvit/canvitb16-abl-rw-stride6-2026-03-03",
    probe_repo="canvit/probe-ade20k-40k-s512-c32-abl-rw-stride6",
).eval()
scene = preprocess(512)(Image.open("scene.jpg").convert("RGB")).unsqueeze(0)  # [1, 3, 512, 512]

state = model.init_state(batch_size=1, canvas_grid_size=32)
with torch.inference_mode():
    viewpoint = Viewpoint.full_scene(batch_size=1, device=scene.device)
    glimpse = sample_at_viewpoint(spatial=scene, viewpoint=viewpoint, glimpse_size_px=128)
    logits, state = model(glimpse=glimpse, state=state, viewpoint=viewpoint)  # [1, 150, 32, 32]
labels = logits.argmax(dim=1)  # ADE20K classes, named in CLASS_NAMES

Details


Features	the 32 × 32 canvas of canvit/canvitb16-abl-rw-stride6-2026-03-03
Probe	LayerNorm, dropout, BatchNorm, 1 × 1 convolution
Training rollouts	10 glimpses of 128 px, R-IID viewpoints, 512 px scenes
Training steps	40,000, batch size 16
Released weights	step 34,000, selected on ADE20K validation mIoU
Optimizer	AdamW, peak learning rate 0.0003, weight decay 0.001
Schedule	1,500-step linear warmup, then cosine decay
Augmentation	DINOv3's segmentation training transforms: scale jitter by a factor of 0.5 to 2, random 512 px crops with its class-ratio crop rejection, horizontal flips with probability 0.5, photometric distortion
Dropout	0.1
Precision	bfloat16 autocast

License

MIT. CanViT is pretrained by distillation from DINOv3, which is released under the DINOv3 License.

Citation

@inproceedings{berreby2026canvit,
  title={CanViT: Toward Active-Vision Foundation Models},
  author={Berreby, Yoha{\"i}-Eliel and Du, Sabrina and Durand, Audrey and Krishna, B. Suresh},
  booktitle={Advances in Neural Information Processing Systems},
  year={2026},
  eprint={2603.22570},
  archivePrefix={arXiv},
  primaryClass={cs.CV},
  url={https://arxiv.org/abs/2603.22570}
}

canvit-pytorch 0.1

This repository's files for canvit-pytorch 0.1 remain at revision canvit-pytorch-0.1: with canvit-pytorch<0.2, pass revision="canvit-pytorch-0.1" to from_pretrained.