autograd-gpu

A neural network framework built from scratch. Implements automatic differentiation, GPU-accelerated operations, and a full suite of layers, optimizers, and data loaders. No PyTorch, no TensorFlow.

Backends: NumPy (CPU), CuPy (CUDA), MLX (Apple Metal). The backend is selected at startup and the same code runs on all three.

---

What is here

autograd_gpu/
  backend/       -- backend abstraction (NumPy / CuPy / MLX)
  tensor.py      -- Tensor class with autograd graph tracking
  functional/
    ops.py       -- differentiable primitives (add, mul, matmul, sum, ...)
    activations.py -- relu, sigmoid, tanh, softmax, gelu, leaky_relu
    loss.py      -- mse, cross_entropy, binary_cross_entropy, huber
  nn/
    module.py    -- Module base class (parameter registration, zero_grad)
    linear.py    -- Linear (fully connected) layer
    conv.py      -- Conv2d, MaxPool2d, AvgPool2d (im2col implementation)
    normalization.py -- BatchNorm1d, BatchNorm2d, LayerNorm
    dropout.py   -- Dropout (inverted, train/eval aware)
    recurrent.py -- RNN, LSTM, GRU (unrolled BPTT)
    container.py -- Sequential
    embedding.py -- Embedding (lookup table)
  optim/
    sgd.py       -- SGD with momentum, Nesterov, weight decay
    adam.py      -- Adam, AdamW (decoupled weight decay)
    schedulers.py -- StepLR, CosineAnnealingLR, ReduceLROnPlateau, OneCycleLR
  data/
    dataset.py   -- Dataset base class
    dataloader.py -- DataLoader with shuffling and batching
    datasets.py  -- MNIST, CIFAR-10 (downloads on first use)
demos/
  mnist.py       -- CNN demo, ~99% test accuracy in 12 epochs on GPU
  cifar10.py     -- MiniResNet demo, ~85% test accuracy in 50 epochs on GPU
docs/
  architecture.md -- design decisions: graph model, im2col, BatchNorm backward
  gpu-setup.md   -- CUDA and Metal installation and troubleshooting

---

Install

pip install -e .                # CPU only (NumPy)
pip install -e ".[cuda]"        # + CUDA (CuPy)
pip install -e ".[metal]"       # + Metal (MLX, Apple Silicon only)
pip install -e ".[dev,demos]"   # + test tools and demo dependencies

---

Quick start

import autograd_gpu as ag
from autograd_gpu import nn, optim

# Backend auto-detected (CUDA > Metal > NumPy)
print(ag.get_backend().name)

# Build a model
model = nn.Sequential(
    nn.Linear(784, 256),
    nn.Dropout(0.2),
    nn.Linear(256, 10),
)

optimizer = optim.Adam(list(model.parameters()), lr=1e-3)

# Training step
x = ag.randn(32, 784, requires_grad=False)
y = ag.tensor([0] * 32)

optimizer.zero_grad()
logits = model(x)
loss = nn.cross_entropy_loss(logits, y)
loss.backward()
optimizer.step()

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

---

Autograd example

The backward pass is triggered by calling .backward() on a scalar Tensor. Gradients accumulate in .grad on all tensors with requires_grad=True.

import autograd_gpu as ag
from autograd_gpu.functional import ops

ag.set_backend("numpy")

x = ag.tensor([2.0, 3.0], requires_grad=True)
y = (x * x).sum()   # y = x0^2 + x1^2
y.backward()
print(x.grad)        # [4.0, 6.0]  (dy/dx = 2x)

---

Running demos

# MNIST -- auto-detect backend
python demos/mnist.py

# MNIST -- force CUDA
python demos/mnist.py --backend cuda --epochs 12

# CIFAR-10 -- Apple Metal (recommended for M-series Macs)
python demos/cifar10.py --backend metal --epochs 50

---

Running tests

pytest

Gradient correctness is verified against finite-difference approximations. Layer tests verify output shapes, gradient flow, and train/eval behavior. Optimizer tests verify convergence on a simple quadratic.

---

Design notes

The engine uses define-by-run (dynamic graph), the same model as PyTorch. Each operation builds a node in the graph during the forward pass. backward() does a reverse topological sort and calls each node's closure.

Conv2d is implemented via im2col: input patches are rearranged into a 2D matrix so convolution becomes a single matmul. This is efficient on GPU because the matmul runs as a single cuBLAS call.

BatchNorm backward is derived analytically rather than through the autograd graph, following the derivation in Ioffe & Szegedy (2015).

LSTM and GRU are unrolled in Python. For very long sequences, truncated BPTT (detaching the hidden state every K steps) prevents memory exhaustion.

See docs/architecture.md for full design discussion. See docs/gpu-setup.md for CUDA and Metal installation.

---

References

• Rumelhart, Hinton & Williams (1986): Learning representations by back-propagating errors
• Hochreiter & Schmidhuber (1997): Long Short-Term Memory
• Ioffe & Szegedy (2015): Batch Normalization
• Kingma & Ba (2015): Adam -- A Method for Stochastic Optimization
• Loshchilov & Hutter (2019): Decoupled Weight Decay Regularization (AdamW)
• Smith & Touvron (2018): A disciplined approach to neural network hyper-parameters (1cycle)
• He et al. (2015): Delving deep into rectifiers (Kaiming initialization)