Skip to content

Models

Vision models for tiny-pytorch implementation.

This module provides pre-built neural network architectures for computer vision tasks, specifically designed for image classification. The models are built using the core neural network components from the tiny-pytorch framework.

The module includes implementations of popular architectures adapted for the tiny-pytorch ecosystem, focusing on efficiency and educational value while maintaining compatibility with the framework's tensor operations and automatic differentiation system.

Key Features
  • Pre-built vision models for image classification
  • Residual network architectures with skip connections
  • Efficient implementations optimized for the tiny-pytorch framework
  • Educational models that demonstrate modern CNN design patterns

Classes:

  • ResNet9 –

    A lightweight ResNet architecture designed for efficient training and inference on smaller datasets. Features residual connections and progressive channel expansion for 10-class classification tasks.

  • ResidualBlock –

    A utility function that creates residual blocks with two ConvBN layers and skip connections, essential for building deep residual networks.

Notes

All models in this module are designed to work with the tiny-pytorch tensor system and support automatic differentiation. Input tensors should be in NCHW format (batch, channels, height, width) and models output classification logits.

Examples:

>>> from tiny_pytorch.vision.models import ResNet9
>>>
>>> # Create a ResNet-9 model for image classification
>>> model = ResNet9()
>>>
>>> # Prepare input data (batch_size=32, channels=3, height=32, width=32)
>>> x = Tensor.randn(32, 3, 32, 32)
>>>
>>> # Forward pass to get classification logits
>>> output = model(x)  # shape: (32, 10)
>>>
>>> # The model is ready for training with appropriate loss functions
>>> # and optimizers from the tiny-pytorch framework

ResNet9

Bases: Module

ResNet-9: A lightweight ResNet architecture for image classification.

ResNet-9 is a simplified version of the ResNet architecture designed for efficient training and inference on smaller datasets. It consists of convolutional layers with residual connections, followed by fully connected layers for classification.

The architecture follows this pattern: 1. Initial convolution (7x7, stride 4) to reduce spatial dimensions 2. Multiple convolutional blocks with residual connections 3. Global average pooling (via Flatten) 4. Fully connected layers for classification

Parameters:

  • device (Device, default: None ) –

    Device on which to place the model parameters. Default is None (uses default device).

Attributes:

  • model (Sequential) –

    The complete ResNet-9 model as a sequential container.

Notes
  • Input is expected to be in NCHW format (batch, channels, height, width).
  • The model is designed for 10-class classification (e.g., CIFAR-10).
  • Uses residual connections to help with gradient flow in deeper layers.
  • The architecture progressively reduces spatial dimensions while increasing the number of channels.

Examples:

>>> model = ResNet9()
>>> x = Tensor.randn(32, 3, 32, 32)  # batch_size=32, channels=3, height=32, width=32
>>> output = model(x)  # shape: (32, 10) - 10 class probabilities

__init__(device=None)

Initialize the ResNet-9 model.

Parameters:

  • device (Device, default: None ) –

    Device on which to place the model parameters. Default is None (uses default device).

forward(x)

Forward pass of the ResNet-9 model.

Parameters:

  • x (Tensor) –

    Input tensor of shape (batch_size, 3, height, width) in NCHW format. Typically used with 32x32 or 64x64 images.

Returns:

  • Tensor –

    Output tensor of shape (batch_size, 10) containing class logits for 10-class classification.

ResidualBlock(in_channels, out_channels, kernel_size, stride, device=None)

Create a residual block with two ConvBN layers and a skip connection.

This function creates a residual block that consists of two ConvBN layers followed by a residual connection that adds the input to the output. The residual connection helps with gradient flow and enables training of deeper networks.

Parameters:

  • in_channels (int) –

    Number of input channels.

  • out_channels (int) –

    Number of output channels.

  • kernel_size (int) –

    Size of the convolutional kernel.

  • stride (int) –

    Stride of the convolution.

  • device (Device, default: None ) –

    Device on which to place the parameters. Default is None (uses default device).

Returns:

  • Residual –

    A residual block module that applies two ConvBN layers with a skip connection.

Notes

The residual block applies two ConvBN operations in sequence, then adds the original input to the result. This helps with gradient flow in deep networks.