NNRT: Neural Network Runtime

NNRT is a lightweight deep learning framework built from scratch in Python.
It is designed to replicate core concepts of modern frameworks like PyTorch, including tensors, autograd, neural networks, and optimizers.

Features

NNRT provides core deep learning components including:

Tensor Core

• Multi dimensional Tensor support
• Autograd (automatic differentiation)
• Basic math operations
• CPU (and optional CUDA support via nnrt.cuda)

Neural Network Layers

• Linear (Fully Connected)
• Conv2D
• MaxPool2D
• BatchNorm1d
• LayerNorm
• Embedding

Activation Functions

• ReLU
• LeakyReLU
• Sigmoid
• Tanh
• GELU
• Softmax
• LogSoftmax

Model Building

• Sequential API
• Module / Parameter system
• ModuleList support
• Dropout
• Flatten layer

Loss Functions

• MSELoss
• CrossEntropyLoss
• NLLLoss
• BCELoss
• BCEWithLogitsLoss

Optimizers

• SGD
• Adam
• RMSProp

Learning Rate Schedulers

• StepLR
• ExponentialLR
• CosineAnnealingLR

Utilities

• Model save / load
• No-grad context manager

Installation

CPU version

pip install nnrt

CUDA Versions

CUDA11.x

pip install nnrt[cuda11x]

CUDA12.x

pip install nnrt[cuda12x]

Quick Example

import nnrt as nn

model = nn.Sequential(
    nn.Linear(784, 256),
    nn.ReLU(),
    nn.Linear(256, 10)
)

loss_fn = nn.CrossEntropyLoss()
optimizer = nn.Adam(model.parameters(), lr=0.001)

x = nn.randn(32, 784)
y = nn.randint(0, 10, (32,))

out = model(x)
loss = loss_fn(out, y)

loss.backward()
optimizer.step()
optimizer.zero_grad()

Goal

NNRT is built for:

• Learning how deep learning frameworks work internally
• Understanding autograd systems
• Experimenting with custom neural network architectures

Disclaimer

This is an educational framework and not optimized for production-scale workloads.

If you find NNRT useful, consider giving it a star on GitHub!