RepViT

RepViT is based on the "RepViT: Revisiting Mobile CNN From ViT Perspective" paper and its official implementation.

Architecture overview

RepViT adapts mobile CNN blocks using design choices associated with efficient vision transformers. Each block separates spatial token mixing from channel mixing, uses squeeze-excitation selectively, and can fuse its training-time depthwise branches for deployment.

Call model.eval() and then model.reparametrize() before exporting or benchmarking the deployment form.

Paper evidence

These ImageNet-1K results are teacher-distilled scores reported by the authors, not Holocron benchmark results.

Model	Parameters	MACs	Top-1
RepViT-M0.9	5.1M	0.8G	78.7%
RepViT-M1.0	6.8M	1.1G	80.0%
RepViT-M1.1	8.2M	1.3G	80.7%

Controlled Holocron benchmark

The Holocron comparison trains from scratch on Imagenette without a teacher: 176px training crops, 232px resize and 224px validation crops, 20 epochs, effective batch size 32, AMP, AdamP at 1e-3, OneCycle, Mixup 0.2, and label smoothing 0.1. MobileOne-S2 uses the identical command as the baseline.

CUDA measurements remain a separate acceptance gate for issue #499; they are not inferred from local CPU or MPS checks.

Model	Parameters before/after fusion	MACs	Top-1	Top-5	Status
RepViT-M0.9	5,103,560 / 5,067,056	Pending	Pending	Pending	CUDA run required
RepViT-M1.0	6,852,900 / 6,810,312	Pending	Pending	Pending	CUDA run required
RepViT-M1.1	8,288,888 / 8,244,312	Pending	Pending	Pending	CUDA run required
MobileOne-S2	Pending rerun	Pending	Pending	Pending	CUDA run required

No pretrained RepViT checkpoint is published with this implementation.

Model builders

All builders rely on RepViT and accept a custom class count through num_classes.

RepViT

RepViT(channels: list[int], num_blocks: list[int], num_classes: int = 10, in_channels: int = 3)

Bases: Sequential

Implements RepViT as described in "RepViT: Revisiting Mobile CNN From ViT Perspective".

PARAMETER	DESCRIPTION
channels	number of output channels in each stage TYPE: list[int]
num_blocks	number of blocks in each stage TYPE: list[int]
num_classes	number of output classes TYPE: int DEFAULT: 10
in_channels	number of input channels TYPE: int DEFAULT: 3

Source code in holocron/models/classification/repvit.py

def __init__(
    self,
    channels: list[int],
    num_blocks: list[int],
    num_classes: int = 10,
    in_channels: int = 3,
) -> None:
    if len(channels) != 4 or len(num_blocks) != 4:
        raise ValueError("`channels` and `num_blocks` are expected to contain four stages")

    patch_embed = nn.Sequential(
        _ConvNorm(in_channels, channels[0] // 2, 3, stride=2, padding=1),
        nn.GELU(),
        _ConvNorm(channels[0] // 2, channels[0], 3, stride=2, padding=1),
    )
    stages: list[nn.Sequential] = []
    in_planes = channels[0]
    for stage_idx, (out_planes, depth) in enumerate(zip(channels, num_blocks, strict=True)):
        blocks: list[nn.Module] = []
        for block_idx in range(depth):
            stride = 2 if stage_idx > 0 and block_idx == 0 else 1
            # Official configs: SE on the first block, then alternating blocks except stage ends.
            use_se = block_idx == 0 if stage_idx == 0 else block_idx % 2 == 1 and block_idx < depth - 1
            blocks.append(_RepViTBlock(in_planes, out_planes, stride, use_se))
            in_planes = out_planes
        stages.append(nn.Sequential(*blocks))

    super().__init__(
        OrderedDict([
            ("features", nn.Sequential(patch_embed, *stages)),
            ("pool", GlobalAvgPool2d(flatten=True)),
            ("head", _BatchNormLinear(channels[-1], num_classes)),
        ])
    )

reparametrize

reparametrize() -> None

Fuse training-time branches and batch-normalization layers for deployment.

Source code in holocron/models/classification/repvit.py

def reparametrize(self) -> None:
    """Fuse training-time branches and batch-normalization layers for deployment."""
    self.features: nn.Sequential
    patch_embed = cast(nn.Sequential, self.features[0])
    if not isinstance(patch_embed[0], _ConvNorm):
        return
    patch_embed[0] = cast(_ConvNorm, patch_embed[0]).reparametrize()
    patch_embed[-1] = cast(_ConvNorm, patch_embed[-1]).reparametrize()
    for stage in self.features[1:]:
        for block in cast(nn.Sequential, stage):
            cast(_RepViTBlock, block).reparametrize()
    self.head = cast(_BatchNormLinear, self.head).reparametrize()

repvit_m0_9

repvit_m0_9(pretrained: bool = False, checkpoint: Checkpoint | None = None, progress: bool = True, **kwargs: Any) -> RepViT

RepViT-M0.9 model.

PARAMETER	DESCRIPTION
pretrained	If True, loads the default checkpoint when one is available TYPE: bool DEFAULT: False
checkpoint	If specified, sets the model parameters to the checkpoint values TYPE: Checkpoint | None DEFAULT: None
progress	If True, displays a download progress bar TYPE: bool DEFAULT: True
kwargs	keyword arguments of RepViT TYPE: Any DEFAULT: {}

RETURNS	DESCRIPTION
RepViT	A RepViT-M0.9 model

Source code in holocron/models/classification/repvit.py

def repvit_m0_9(
    pretrained: bool = False,
    checkpoint: Checkpoint | None = None,
    progress: bool = True,
    **kwargs: Any,
) -> RepViT:
    """RepViT-M0.9 model.

    Args:
        pretrained: If True, loads the default checkpoint when one is available
        checkpoint: If specified, sets the model parameters to the checkpoint values
        progress: If True, displays a download progress bar
        kwargs: keyword arguments of [`RepViT`][holocron.models.classification.repvit.RepViT]

    Returns:
        A RepViT-M0.9 model
    """
    checkpoint = _handle_legacy_pretrained(pretrained, checkpoint, None)
    return _repvit(checkpoint, progress, [48, 96, 192, 384], [3, 4, 16, 3], **kwargs)

repvit_m1_0

repvit_m1_0(pretrained: bool = False, checkpoint: Checkpoint | None = None, progress: bool = True, **kwargs: Any) -> RepViT

RepViT-M1.0 model.

PARAMETER	DESCRIPTION
pretrained	If True, loads the default checkpoint when one is available TYPE: bool DEFAULT: False
checkpoint	If specified, sets the model parameters to the checkpoint values TYPE: Checkpoint | None DEFAULT: None
progress	If True, displays a download progress bar TYPE: bool DEFAULT: True
kwargs	keyword arguments of RepViT TYPE: Any DEFAULT: {}

RETURNS	DESCRIPTION
RepViT	A RepViT-M1.0 model

Source code in holocron/models/classification/repvit.py

def repvit_m1_0(
    pretrained: bool = False,
    checkpoint: Checkpoint | None = None,
    progress: bool = True,
    **kwargs: Any,
) -> RepViT:
    """RepViT-M1.0 model.

    Args:
        pretrained: If True, loads the default checkpoint when one is available
        checkpoint: If specified, sets the model parameters to the checkpoint values
        progress: If True, displays a download progress bar
        kwargs: keyword arguments of [`RepViT`][holocron.models.classification.repvit.RepViT]

    Returns:
        A RepViT-M1.0 model
    """
    checkpoint = _handle_legacy_pretrained(pretrained, checkpoint, None)
    return _repvit(checkpoint, progress, [56, 112, 224, 448], [3, 4, 16, 3], **kwargs)

repvit_m1_1

repvit_m1_1(pretrained: bool = False, checkpoint: Checkpoint | None = None, progress: bool = True, **kwargs: Any) -> RepViT

RepViT-M1.1 model.

PARAMETER	DESCRIPTION
pretrained	If True, loads the default checkpoint when one is available TYPE: bool DEFAULT: False
checkpoint	If specified, sets the model parameters to the checkpoint values TYPE: Checkpoint | None DEFAULT: None
progress	If True, displays a download progress bar TYPE: bool DEFAULT: True
kwargs	keyword arguments of RepViT TYPE: Any DEFAULT: {}

RETURNS	DESCRIPTION
RepViT	A RepViT-M1.1 model

Source code in holocron/models/classification/repvit.py

def repvit_m1_1(
    pretrained: bool = False,
    checkpoint: Checkpoint | None = None,
    progress: bool = True,
    **kwargs: Any,
) -> RepViT:
    """RepViT-M1.1 model.

    Args:
        pretrained: If True, loads the default checkpoint when one is available
        checkpoint: If specified, sets the model parameters to the checkpoint values
        progress: If True, displays a download progress bar
        kwargs: keyword arguments of [`RepViT`][holocron.models.classification.repvit.RepViT]

    Returns:
        A RepViT-M1.1 model
    """
    checkpoint = _handle_legacy_pretrained(pretrained, checkpoint, None)
    return _repvit(checkpoint, progress, [64, 128, 256, 512], [3, 4, 14, 3], **kwargs)