computer_vision.vision_transformer
==================================

.. py:module:: computer_vision.vision_transformer

.. autoapi-nested-parse::

   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
----------

.. autoapisummary::

   computer_vision.vision_transformer.rng


Functions
---------

.. autoapisummary::

   computer_vision.vision_transformer.add_positional_encoding
   computer_vision.vision_transformer.attention_mechanism
   computer_vision.vision_transformer.create_patches
   computer_vision.vision_transformer.feedforward_network
   computer_vision.vision_transformer.layer_norm
   computer_vision.vision_transformer.patch_embedding
   computer_vision.vision_transformer.transformer_encoder_block
   computer_vision.vision_transformer.vision_transformer


Module Contents
---------------

.. py:function:: 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)


.. py:function:: 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


.. py:function:: 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)


.. py:function:: 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)


.. py:function:: 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


.. py:function:: 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)


.. py:function:: 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)


.. py:function:: 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,)


.. py:data:: rng