See our collection for all versions of MiT.

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

zeromodels/mit_b5_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-b5 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
import numpy as np
from zeromodels.models.mit import MiTImageClassify, MiTModel

model = MiTImageClassify.from_weights("zeromodels/mit_b5_in1k")
backbone = MiTModel.from_weights(
    "zeromodels/mit_b5_in1k", as_backbone=True
)

image = Image.open("your_image.jpg").convert("RGB")
image = image.resize((224, 224))
x = np.asarray(image, dtype="float32")[None]  # (1, H, W, 3)
print(model(x).shape)  # (1, num_classes)
feats = backbone(x)
print(len(feats), [tuple(f.shape) for f in feats])

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

Variant	Hub
`mit_b0_in1k`	`zeromodels/mit_b0_in1k`
`mit_b1_in1k`	`zeromodels/mit_b1_in1k`
`mit_b2_in1k`	`zeromodels/mit_b2_in1k`
`mit_b3_in1k`	`zeromodels/mit_b3_in1k`
`mit_b4_in1k`	`zeromodels/mit_b4_in1k`
`mit_b5_in1k`	`zeromodels/mit_b5_in1k`

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-b5").