import torch
from canvit_pytorch import CanViTForImageClassification, Viewpoint, sample_at_viewpoint
from canvit_pytorch.preprocess import preprocess
from PIL import Image

clf = CanViTForImageClassification.from_pretrained(
    "canvit/canvitb16-add-vpe-finetune-g128px-s512px-in1k-2026-04-06"
).eval()
assert isinstance(clf, CanViTForImageClassification)
assert clf.n_classes == 1000
assert clf.local_dim == 768

image = preprocess(512)(Image.open("cat.jpg").convert("RGB"))
assert isinstance(image, torch.Tensor)
image = image.unsqueeze(0)  # [1, 3, 512, 512]

state = clf.init_state(batch_size=1, canvas_grid_size=32)

# Process one or more glimpses sequentially
with torch.inference_mode():
    vp = Viewpoint.full_scene(batch_size=1, device=image.device)
    glimpse = sample_at_viewpoint(spatial=image, viewpoint=vp, glimpse_size_px=128)
    logits, state = clf(glimpse=glimpse, state=state, viewpoint=vp)

assert logits.shape == (1, 1000)
pred = logits.argmax(dim=-1)  # ImageNet-1K class index

Training

Base checkpoint: canvit/canvitb16-add-vpe-pretrain-g128px-s512px-in21k-dv3b16-2026-02-02
Head init: Pretrained CLS projection, destandardization, and DINOv3 linear probe (yberreby/dinov3-vitb16-lvd1689m-in1k-512x512-linear-clf-probe) algebraically fused into a single LN → Linear head
Data: ImageNet-1K training set
Augmentation: torchvision.transforms.v2.RandomResizedCrop(512, scale=(0.2, 1.0)), RandomHorizontalFlip
Episodes: 4 F-IID glimpses per scene — \( t_0 \) = full scene, \( t_{>0} \) = random with \( p(s) \propto (1-s) \), min_scale=0.05 (matching pretraining). Full BPTT
Loss: Cross-entropy at every timestep, averaged over all glimpses in the trajectory. Label smoothing 0.1
Optimizer: AdamW — \( \beta_1 = 0.9 \), \( \beta_2 = 0.999 \), \( \varepsilon = 10^{-8} \). LR 2.5e-5, WD 1e-4, gradient clipping max norm 1.0
LR schedule: Linear warmup from 0 to peak LR over 25,000 steps, then cosine decay to 0
Batch size: 256 scenes (1,024 glimpses/step)
Duration: 20 epochs (100,080 steps)
All parameters trainable (backbone + head)
Software: PyTorch 2.9.0, torch_xla 2.9.0
Hardware: Google Cloud TPU v6e-4 (4 chips, SPMD)

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}
}