canvit/canvitb16-add-vpe-finetune-g128px-s512px-in1k-2026-04-06-nnx

CanViT-B/16 Finetuned IN1K (JAX / Flax NNX)

JAX-native checkpoint for CanViT finetuned on ImageNet-1K, converted from the PyTorch checkpoint.

Usage

uv add "canvit-nnx @ git+https://github.com/yberreby/CanViT-NNX.git"

from canvit_nnx import CanViTForImageClassification, Viewpoint, sample_at_viewpoint

clf = CanViTForImageClassification.from_pretrained("canvit/canvitb16-add-vpe-finetune-g128px-s512px-in1k-2026-04-06-nnx")
state = clf.init_state(batch_size=1, canvas_grid_size=32)

vp = Viewpoint.full_scene(batch_size=1)
glimpse = sample_at_viewpoint(spatial=image, viewpoint=vp, glimpse_size_px=128)
logits, state = clf(glimpse, state, vp)

Source: CanViT-NNX

Citation

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