computer_vision.vision_transformer

Vision Transformer (ViT) for Image Classification

This module implements the Vision Transformer architecture as described in the paper “An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale” by Dosovitskiy et al. (2020)

Paper: https://arxiv.org/abs/2010.11929

The Vision Transformer splits an image into fixed-size patches, linearly embeds each patch, adds position embeddings, and feeds the resulting sequence of vectors to a standard Transformer encoder. For classification, a learnable classification token is prepended to the sequence.

Author: devvratpathak

Attributes

rng

Functions

add_positional_encoding(→ numpy.ndarray)	Add learnable positional encodings to patch embeddings.
attention_mechanism(→ tuple[numpy.ndarray, numpy.ndarray])	Compute scaled dot-product attention.
create_patches(→ tuple[numpy.ndarray, tuple[int, int]])	Split an image into non-overlapping patches.
feedforward_network(→ numpy.ndarray)	Apply position-wise feed-forward network.
layer_norm(→ numpy.ndarray)	Apply Layer Normalization.
patch_embedding(→ numpy.ndarray)	Linearly project flattened patches to embedding dimension.
transformer_encoder_block(→ numpy.ndarray)	Apply a single Transformer encoder block.
vision_transformer(→ numpy.ndarray)	Apply Vision Transformer for image classification.

Module Contents

computer_vision.vision_transformer.add_positional_encoding(embeddings: numpy.ndarray, num_positions: int | None = None) → numpy.ndarray

Add learnable positional encodings to patch embeddings.

Args:

embeddings: Embedded patches of shape (num_patches, embedding_dim) num_positions: Number of positions (if None, uses num_patches + 1 for CLS token)

Returns:

Embeddings with positional encoding of shape (num_positions, embedding_dim)

Examples:

>>> embeddings = np.random.rand(4, 768)
>>> pos_embeddings = add_positional_encoding(embeddings)
>>> pos_embeddings.shape
(5, 768)

>>> embeddings = np.random.rand(196, 512)
>>> pos_embeddings = add_positional_encoding(embeddings)
>>> pos_embeddings.shape
(197, 512)

computer_vision.vision_transformer.attention_mechanism(query: numpy.ndarray, key: numpy.ndarray, value: numpy.ndarray, mask: numpy.ndarray | None = None) → tuple[numpy.ndarray, numpy.ndarray]

Compute scaled dot-product attention.

Attention(Q, K, V) = softmax(QK^T / sqrt(d_k))V

Args:

query: Query matrix of shape (seq_len, d_k) key: Key matrix of shape (seq_len, d_k) value: Value matrix of shape (seq_len, d_v) mask: Optional attention mask

Returns:

A tuple containing: - output: Attention output of shape (seq_len, d_v) - attention_weights: Attention weights of shape (seq_len, seq_len)

Examples:

>>> q = np.random.rand(10, 64)
>>> k = np.random.rand(10, 64)
>>> v = np.random.rand(10, 64)
>>> output, weights = attention_mechanism(q, k, v)
>>> output.shape
(10, 64)
>>> weights.shape
(10, 10)
>>> np.allclose(weights.sum(axis=1), 1.0)
True

computer_vision.vision_transformer.create_patches(image: numpy.ndarray, patch_size: int = 16) → tuple[numpy.ndarray, tuple[int, int]]

Split an image into non-overlapping patches.

Args:

image: Input image array of shape (height, width, channels) patch_size: Size of each square patch (default: 16)

Returns:

A tuple containing: - patches: Array of shape (num_patches, patch_size, patch_size, channels) - grid_size: Tuple (height_patches, width_patches) representing the grid

Examples:

>>> img = np.random.rand(32, 32, 3)
>>> patches, grid = create_patches(img, patch_size=16)
>>> patches.shape
(4, 16, 16, 3)
>>> grid
(2, 2)

>>> img = np.random.rand(224, 224, 3)
>>> patches, grid = create_patches(img, patch_size=16)
>>> patches.shape
(196, 16, 16, 3)
>>> grid
(14, 14)

computer_vision.vision_transformer.feedforward_network(embeddings: numpy.ndarray, hidden_dim: int = 3072) → numpy.ndarray

Apply position-wise feed-forward network.

FFN(x) = max(0, xW1 + b1)W2 + b2

Args:

embeddings: Input array of shape (seq_len, embedding_dim) hidden_dim: Hidden dimension size (default: 3072, typically 4x embedding_dim)

Returns:

Output array of shape (seq_len, embedding_dim)

Examples:

>>> embeddings = np.random.rand(10, 768)
>>> output = feedforward_network(embeddings, hidden_dim=3072)
>>> output.shape
(10, 768)

>>> embeddings = np.random.rand(197, 512)
>>> output = feedforward_network(embeddings, hidden_dim=2048)
>>> output.shape
(197, 512)

computer_vision.vision_transformer.layer_norm(embeddings: numpy.ndarray, epsilon: float = 1e-06) → numpy.ndarray

Apply Layer Normalization.

Args:

embeddings: Input array of shape (seq_len, embedding_dim) epsilon: Small constant for numerical stability (default: 1e-6)

Returns:

Normalized array of same shape as input

Examples:

>>> embeddings = np.random.rand(10, 768)
>>> normalized = layer_norm(embeddings)
>>> normalized.shape
(10, 768)
>>> np.allclose(normalized.mean(axis=1), 0.0, atol=1e-6)
True
>>> np.allclose(normalized.std(axis=1), 1.0, atol=1e-6)
True

computer_vision.vision_transformer.patch_embedding(patches: numpy.ndarray, embedding_dim: int = 768) → numpy.ndarray

Linearly project flattened patches to embedding dimension.

Args:

patches: Array of patches with shape

(num_patches, patch_size, patch_size, channels)

embedding_dim: Dimension of the embedding space (default: 768)

Returns:

Embedded patches of shape (num_patches, embedding_dim)

Examples:

>>> patches = np.random.rand(4, 16, 16, 3)
>>> embeddings = patch_embedding(patches, embedding_dim=768)
>>> embeddings.shape
(4, 768)

>>> patches = np.random.rand(196, 16, 16, 3)
>>> embeddings = patch_embedding(patches, embedding_dim=512)
>>> embeddings.shape
(196, 512)

computer_vision.vision_transformer.transformer_encoder_block(embeddings: numpy.ndarray, num_heads: int = 12, hidden_dim: int = 3072) → numpy.ndarray

Apply a single Transformer encoder block.

The block consists of: 1. Multi-head self-attention with residual connection and layer norm 2. Feed-forward network with residual connection and layer norm

Args:

embeddings: Input array of shape (seq_len, embedding_dim) num_heads: Number of attention heads (default: 12, kept for API) hidden_dim: Hidden dimension for FFN (default: 3072)

Returns:

Output array of shape (seq_len, embedding_dim)

Examples:

>>> embeddings = np.random.rand(197, 768)
>>> output = transformer_encoder_block(
...     embeddings, num_heads=12, hidden_dim=3072
... )
>>> output.shape
(197, 768)

>>> embeddings = np.random.rand(50, 512)
>>> output = transformer_encoder_block(
...     embeddings, num_heads=8, hidden_dim=2048
... )
>>> output.shape
(50, 512)

computer_vision.vision_transformer.vision_transformer(image: numpy.ndarray, patch_size: int = 16, embedding_dim: int = 768, num_layers: int = 12, num_heads: int = 12, hidden_dim: int = 3072, num_classes: int = 1000) → numpy.ndarray

Apply Vision Transformer for image classification.

Architecture: 1. Split image into patches 2. Linear projection of flattened patches 3. Add positional embeddings 4. Pass through Transformer encoder layers 5. Extract CLS token and apply classification head

Args:

image: Input image array of shape (height, width, channels) patch_size: Size of each patch (default: 16) embedding_dim: Embedding dimension (default: 768) num_layers: Number of Transformer layers (default: 12) num_heads: Number of attention heads (default: 12) hidden_dim: Hidden dimension in FFN (default: 3072) num_classes: Number of output classes (default: 1000)

Returns:

Class logits of shape (num_classes,)

Examples:

>>> img = np.random.rand(224, 224, 3)
>>> logits = vision_transformer(img, patch_size=16, num_classes=10)
>>> logits.shape
(10,)

>>> img = np.random.rand(32, 32, 3)
>>> logits = vision_transformer(
...     img, patch_size=16, embedding_dim=512, num_layers=6, num_classes=100
... )
>>> logits.shape
(100,)

computer_vision.vision_transformer.rng