Implementation:Kornia Kornia MobileViT

Knowledge Sources	Kornia
Domains	Vision, Mobile_Networks, Transformer
Last Updated	2026-02-09 15:00 GMT

Overview

MobileViT implements a lightweight hybrid architecture combining MobileNetV2 blocks with Vision Transformer blocks for efficient mobile deployment, as described in the paper MobileViT: Light-weight, General-purpose, and Mobile-friendly Vision Transformer (arXiv:2110.02178).

Description

This module provides the MobileViT class and its building blocks within the Kornia library. The architecture interleaves MV2Block (MobileNetV2 inverted residual bottleneck) layers for local feature extraction with MobileViTBlock layers that combine local convolutions with global self-attention via a Transformer. The MobileViTBlock performs a fold-unfold operation to convert spatial features into sequences for transformer processing. The model supports three size variants: xxs (extra-extra-small), xs (extra-small), and s (small). Helper modules include PreNorm, FeedForward, Attention, Transformer, and utility functions conv_1x1_bn and conv_nxn_bn.

Usage

Import this module when you need a mobile-optimized vision backbone that balances local and global feature processing. The output is a feature map suitable for downstream tasks like classification (with an additional head) or dense prediction.

Code Reference

Source Location

• Repository: Kornia
• File: kornia/models/vit_mobile.py
• Lines: 1-342

Signature

class MobileViT(nn.Module):
    def __init__(
        self,
        mode: str = "xxs",
        in_channels: int = 3,
        patch_size: Tuple[int, int] = (2, 2),
        dropout: float = 0.0,
    ) -> None: ...

    def forward(self, x: torch.Tensor) -> torch.Tensor: ...

Import

from kornia.models.vit_mobile import MobileViT

I/O Contract

Inputs

Name	Type	Required	Description
mode	str	No	Model size variant: xxs, xs, or s (default xxs).
in_channels	int	No	Number of input image channels (default 3).
patch_size	Tuple[int, int]	No	Patch size for folding/unfolding in MobileViTBlock (default (2, 2)). Image size must be divisible by this.
dropout	float	No	Dropout rate for transformer layers (default 0.0).

Outputs

Name	Type	Description
output	torch.Tensor	Feature map of shape (B, C_out, H/32, W/32) where C_out depends on the mode: 320 for xxs, 384 for xs, 640 for s.

Key Components

MV2Block

MobileNetV2 inverted residual block with expansion, depthwise separable convolution, and optional residual connection. Uses SiLU activation.

MobileViTBlock

Hybrid block that applies local convolutions, reshapes spatial features into a patch sequence, processes them with a Transformer, and fuses the result with a skip connection.

Transformer

Standard transformer with PreNorm, Attention, and FeedForward layers. Uses 4 attention heads with head dimension 8.

Attention

Multi-head attention with QKV projection, scaled dot-product attention, and optional output projection.

Model Configurations

Mode	Expansion	Dims	Output Channels
xxs	2	[64, 80, 96]	320
xs	4	[96, 120, 144]	384
s	4	[144, 192, 240]	640

Usage Examples

import torch
from kornia.models.vit_mobile import MobileViT

# Create MobileViT-XXS
img = torch.rand(1, 3, 256, 256)
mvit = MobileViT(mode='xxs')
output = mvit(img)  # shape: (1, 320, 8, 8)

# Create MobileViT-S
mvit_s = MobileViT(mode='s')
output_s = mvit_s(img)  # shape: (1, 640, 8, 8)

Related Pages

• Principle:Kornia_Kornia_Mobile_Vision_Transformer