Implementation:Kornia Kornia TinyViT

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

Overview

TinyViT implements a compact Vision Transformer architecture that combines convolutional layers with self-attention for efficient image classification and feature extraction, as described in the paper TinyViT: Fast Pretraining Distillation for Small Vision Transformers (arXiv:2207.10666).

Description

This module provides the TinyViT model and all of its building blocks within the Kornia library. TinyViT is designed to be a parameter-efficient vision transformer that uses Mobile Inverted Bottleneck Convolutions (MBConv) in early stages and windowed self-attention in later stages. The architecture supports three model sizes (5M, 11M, 21M parameters) and can optionally be configured for MobileSAM feature extraction with mobile_sam=True. Key submodules include ConvBN, PatchEmbed, MBConv, PatchMerging, ConvLayer, MLP, Attention, TinyViTBlock, and BasicLayer.

Usage

Import this module when you need a lightweight vision transformer for image classification on ImageNet or as a backbone for the MobileSAM segmentation pipeline. Use TinyViT.from_config() for convenient instantiation of pre-defined model variants with optional pre-trained weights.

Code Reference

Source Location

• Repository: Kornia
• File: kornia/models/tiny_vit.py
• Lines: 1-691

Signature

class TinyViT(nn.Module):
    def __init__(
        self,
        img_size: int = 224,
        in_chans: int = 3,
        num_classes: int = 1000,
        embed_dims: Sequence[int] = (96, 192, 384, 768),
        depths: Sequence[int] = (2, 2, 6, 2),
        num_heads: Sequence[int] = (3, 6, 12, 24),
        window_sizes: Sequence[int] = (7, 7, 14, 7),
        mlp_ratio: float = 4.0,
        drop_rate: float = 0.0,
        drop_path_rate: float = 0.0,
        use_checkpoint: bool = False,
        mbconv_expand_ratio: float = 4.0,
        local_conv_size: int = 3,
        activation: type[nn.Module] = nn.GELU,
        mobile_sam: bool = False,
    ) -> None: ...

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

    @staticmethod
    def from_config(variant: str, pretrained: bool | str = False, **kwargs: Any) -> TinyViT: ...

Import

from kornia.models.tiny_vit import TinyViT

I/O Contract

Inputs

Name	Type	Required	Description
img_size	int	No	Size of the input image (default 224).
in_chans	int	No	Number of input channels (default 3).
num_classes	int	No	Number of output classes for classification head (default 1000).
embed_dims	Sequence[int]	No	Embedding dimensions for each stage (default (96, 192, 384, 768)).
depths	Sequence[int]	No	Number of blocks per stage (default (2, 2, 6, 2)).
num_heads	Sequence[int]	No	Number of attention heads per stage (default (3, 6, 12, 24)).
window_sizes	Sequence[int]	No	Window sizes for local attention per stage (default (7, 7, 14, 7)).
mlp_ratio	float	No	Ratio of MLP hidden dim to embedding dim (default 4.0).
drop_rate	float	No	Dropout rate (default 0.0).
drop_path_rate	float	No	Stochastic depth rate (default 0.0).
use_checkpoint	bool	No	Whether to use activation checkpointing (default False).
mbconv_expand_ratio	float	No	Expansion ratio for MBConv blocks (default 4.0).
local_conv_size	int	No	Kernel size for local convolution in TinyViTBlock (default 3).
activation	type[nn.Module]	No	Activation function class (default nn.GELU).
mobile_sam	bool	No	Whether to produce feature maps for MobileSAM instead of classification logits (default False).

Outputs

Name	Type	Description
output	torch.Tensor	When mobile_sam=False: classification logits of shape (B, num_classes). When mobile_sam=True: feature maps of shape (B, 256, feat_size, feat_size).

Key Components

ConvBN

A sequential block of Conv2d followed by BatchNorm2d and an optional activation.

PatchEmbed

Converts the input image into patch embeddings via two strided 3x3 convolutions (total stride of 4).

MBConv

Mobile Inverted Residual Bottleneck Convolution with point-wise expansion, depth-wise convolution, and point-wise projection with residual connection and DropPath.

PatchMerging

Downsamples spatial dimensions while increasing channel dimensions using convolutions.

TinyViTBlock

Core block combining windowed multi-head self-attention with a local convolution and a feed-forward MLP, all with residual connections and layer scaling.

Attention

Multi-head attention with separable relative position encoding using decomposed attention biases.

Usage Examples

import torch
from kornia.models.tiny_vit import TinyViT

# Classification mode (default)
model = TinyViT(img_size=224, num_classes=1000)
img = torch.randn(1, 3, 224, 224)
logits = model(img)  # shape: (1, 1000)

# Using from_config with pretrained weights
model = TinyViT.from_config("5m", pretrained="in1k")

# MobileSAM feature extraction mode
model = TinyViT(img_size=1024, mobile_sam=True)
img = torch.randn(1, 3, 1024, 1024)
features = model(img)  # shape: (1, 256, feat_size, feat_size)

Related Pages

• Principle:Kornia_Kornia_Tiny_Vision_Transformer