mit_b2_in1k

See our collection for all versions of MiT.

Run MiT with Keras 3: JAX, PyTorch, or TensorFlow

zeromodels/mit_b2_in1k

Paper: SegFormer: Simple and Efficient Design for Semantic Segmentation with Transformers (arXiv:2105.15203) · HF Papers

MiT is the hierarchical Mix Transformer encoder from SegFormer, also usable for ImageNet classification. For full SegFormer segmentation heads, see the SegFormer collection.

For more details on the model, please go to the upstream model card.

Pure-Keras 3 conversion of nvidia/mit-b2 for zeromodels. One implementation runs unmodified on TensorFlow / Torch / JAX.

This is an image-classification / backbone checkpoint (MiTImageClassify / MiTModel).

✨ Quick start

import os

os.environ["KERAS_BACKEND"] = "torch"  # or "jax" / "tensorflow"

from PIL import Image
from zeromodels.models.mit import MiTImageClassify, MiTModel, MiTImageProcessor

model = MiTImageClassify.from_weights("zeromodels/mit_b2_in1k")
processor = MiTImageProcessor.from_weights("zeromodels/mit_b2_in1k")

image = Image.open("your_image.jpg").convert("RGB")
pixels = processor(image)  # resize + normalize (normalization lives in the processor)
logits = model(pixels, training=False)
print(logits.shape)  # (1, num_classes)

# Feature extraction: the backbone without the classifier head
backbone = MiTModel.from_weights("zeromodels/mit_b2_in1k", as_backbone=True)
features = backbone(pixels, training=False)

Load any MiT variant the same way with from_weights("zeromodels/<variant>"):

Tips

• Set KERAS_BACKEND before importing Keras / zeromodels.
• MiTImageClassify returns class logits; MiTModel returns features (as_backbone=True for multi-scale stages).
• See docs and Loading Weights.
• Upstream / timm checkpoints: MiTImageClassify.from_weights("hf:nvidia/mit-b2").

Special Thanks

A huge thank you to the MiT authors and the timm / Hub communities for creating and releasing these models.

License: see YAML license (usually matches the upstream checkpoint).