Modelpublic

canvit/probe-ade20k-40k-s512-c32-in21k

sourceHugging Facemitupdated 8h agoView on Hugging Face

0likes47downloads

Model Card

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

bash

pip install "canvit-pytorch>=0.2"

python

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-add-vpe-pretrain-g128px-s512px-in21k-dv3b16-2026-02-02",
    probe_repo="canvit/probe-ade20k-40k-s512-c32-in21k",
).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-add-vpe-pretrain-g128px-s512px-in21k-dv3b16-2026-02-02
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
Optimizer	AdamW, peak learning rate 0.0003, weight decay 0.001
Schedule	1,500-step linear warmup, then cosine decay
Augmentation	random crops of scale 0.5 to 2, horizontal flips
Dropout	0.1
Precision	bfloat16 autocast

Citation

bibtex

@article{berreby2026canvit,
  title={CanViT: Toward Active-Vision Foundation Models},
  author={Berreby, Yoha{\"i}-Eliel and Du, Sabrina and Durand, Audrey and Krishna, B. Suresh},
  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.