Tiny-PyTorch 🧠

Unravel the magic of modern deep learning by building a PyTorch-like framework from the ground up.

Tiny-PyTorch is an educational deep learning framework built entirely in Python. It demystifies the core machinery of libraries like PyTorch by providing a clean, focused, and from-scratch implementation of the essential components.

Philosophy: Understanding by Building

The best way to truly understand how complex systems work is to build them yourself. Tiny-PyTorch is born from this philosophy. While production frameworks like PyTorch and TensorFlow provide powerful, high-level abstractions, their internal complexity can be a barrier to learning.

This project strips away those abstractions, allowing you to:

• See the Core Logic: Grasp the fundamental algorithms and data structures that power deep learning, from the Tensor object to the backpropagation process.
• Connect Theory to Code: Bridge the gap between the mathematical concepts of deep learning and their concrete implementation.
• Become a Better Practitioner: Use high-level frameworks more effectively by understanding their internal mechanics, performance trade-offs, and potential pitfalls.

✨ Core Features

• Dynamic Computation Graph: Tensors track their history, allowing for flexible model architectures.
• Reverse-Mode Automatic Differentiation: An efficient gradient calculation engine (autograd) built from scratch.
• Extensible nn.Module System: A familiar API for building complex neural network layers and models.
• Standard Optimizers: Implementations of SGD and Adam to handle parameter updates.
• Hardware Acceleration: A pluggable backend system supporting NumPy, custom CPU (C++), and CUDA (GPU) operations.
• Data Loading Utilities: Dataset and DataLoader classes for efficient data pipelines.

🏗️ Project Architecture

The framework is built in a bottom-up fashion, where each layer of abstraction relies on the one below it. This mirrors the logical structure of major deep learning libraries.

graph TD
    subgraph "High-Level API"
        C[nn.Module] --> D[Optimizers]
        B[Tensors & Autograd] --> C
    end
    subgraph "Low-Level Engine"
        A[NDArray] --> B
        subgraph "Backends"
            direction LR
            np[NumPy]
            cpu[CPU]
            gpu[CUDA]
        end
        Backend --> A
    end

    style C fill:#f9f,stroke:#333,stroke-width:2px
    style B fill:#ccf,stroke:#333,stroke-width:2px
    style A fill:#cfc,stroke:#333,stroke-width:2px

1. Backends (NumPy, CPU, CUDA): Perform the actual mathematical computations on flat arrays of data.
2. NDArray: A generic, strided N-dimensional array class that provides a unified interface over different backends.
3. Tensor & Autograd: The heart of the framework. A Tensor wraps an NDArray and builds a dynamic computation graph. The autograd engine traverses this graph to perform reverse-mode automatic differentiation.
4. High-Level API (nn, optimizer): Provides the familiar modules, layers, and optimization algorithms for building and training neural networks.

🚀 Quick Start

To install Tiny-PyTorch, you have two main options:

Install from PyPI (using pip)

You can install the latest stable version directly from PyPI using pip:

pip install tiny-pytorch

Install from Source (GitHub Repository)

To get the very latest development version or if you plan to contribute, you can install from the GitHub repository:

1. Clone the repository:

   git clone https://github.com/your-username/tiny-pytorch.git
   

2. Navigate to the project directory:

   cd tiny-pytorch
   

3. Install in editable mode: This allows you to make changes to the source code and have them reflected without reinstalling.

   pip install -e .
   

Here's a simple example of defining a model and running a forward/backward pass.

import numpy as np

import tiny_pytorch as tp
import tiny_pytorch.nn as nn
import tiny_pytorch.optim as optim
from tiny_pytorch import init


# 1. Define a simple classifier
class SimpleNet(nn.Module):
    def __init__(self, in_features, num_classes):
        super().__init__()
        self.fc1 = nn.Linear(in_features, 64)
        self.relu = nn.ReLU()
        self.fc2 = nn.Linear(64, num_classes)

    def forward(self, x):
        x = self.fc1(x)
        x = self.relu(x)
        return self.fc2(x)


# 2. Initialize model, optimizer, and loss function
model = SimpleNet(in_features=10, num_classes=3)
optimizer = optim.Adam(model.parameters(), lr=0.001)
loss_fn = nn.SoftmaxLoss()

# 3. Create dummy data
x_train = init.randn(32, 10)
y_true = tp.Tensor(np.random.randint(0, 3, size=(32,)))

# 4. Perform a single training step
optimizer.reset_grad()  # Reset gradients
logits = model(x_train)  # Forward pass
loss = loss_fn(logits, y_true)  # Compute loss
loss.backward()  # Backward pass (autograd)
optimizer.step()  # Update weights

print(f"Loss: {loss.numpy().item():.4f}")

🗺️ Roadmap

The project is developed in two main phases. Our current progress is tracked below.

• Phase I: Core Framework (NumPy Backend)
  • Tensor: The main multi-dimensional array with autograd support.
  • Op: The base class for all tensor operations.
  • Automatic Differentiation: Reverse-mode autograd engine.
  • init: Parameter initialization functions (kaiming, xavier, etc.).
  • nn: Core neural network layers (Linear, ReLU, BatchNorm, Conv2d).
  • optimizer: SGD and Adam optimizers.
  • data: Dataset and DataLoader for data handling.
• Phase II: Hardware Acceleration & Advanced Models
  • NDArray: Generic, strided N-dimensional array.
  • NumPy Backend
  • CPU Backend (C++)
  • CUDA Backend (GPU)
  • Advanced CNN operations (e.g., padding, dilation).
  • ResNet implementation.
  • RNN and LSTM layers.
  • A simple Language Model.

📚 Documentation

The official documentation, including detailed API references and tutorials, is hosted at: https://imaddabbura.github.io/tiny-pytorch/

⚠️ Limitations

As an educational project, Tiny-PyTorch has some intentional simplifications:

• Explicit Broadcasting: Broadcasting for element-wise operations must be done manually if tensor shapes do not match.
• Single Data Type: NDArray only supports the float32 dtype.
• Contiguous Memory: Operations on the underlying 1D array require a call to compact() to ensure the data is in a contiguous memory block.
• Limited Reductions: Reduction operations (e.g., sum, max) can only be performed on a single axis or all axes at once.

🙏 Acknowledgements

This project was inspired by and built upon ideas from the following resources:

• Deep Learning Systems by J. Zico Kolter and Tianqi Chen
• micrograd by Andrej Karpathy
• Machine Learning Systems by Tianqi Chen and Zhihao Jia

License

Tiny-PyTorch is licensed under the Apache License 2.0. See the LICENSE file for details.