Aether-ML

A deep learning framework built from the ground up on NumPy and CuPy, with hand-written GPU kernels for both CUDA and ROCm.

Why Aether-ML?

Aether-ML implements forward/backward propagation, convolutional and fully connected networks, normalization, pooling, and losses without building on top of any existing deep learning library.

With CuPy, the GPU path leverages hand-written CUDA and HIP kernels resolved once at per operation/shape and memoized after. It's an intentionally ground-up project, as I started with no background in ML or GPU kernel programming, and in the process, it helped me understand how deep learning runtimes work under the hood.

All four backends converge to the same loss; the fused-kernel CuPy path is the one built for speed for both GPU vendors. The RTX 3050 Laptops fused kernel path is around 4.6x slower than the 9070 XT's fused kernel path, but still it manages to outpace the 9070 XT's path without custom kernels by over 9x.

Example

import aether as ae

# X_train, y_train, X_test, y_test: CIFAR-10 as CuPy arrays

# --- Preprocessing ---
feature_pipeline = ae.Compose([
    ae.ToTensor(dtype='float32', target_device='cupy'),
    ae.Rescale(factor=1.0 / 255.0),
    ae.StandardScaler()
])

# --- Model: one conv block (32x32 -> 16x16), then a classification head ---
model = ae.Model()
model.manual_seed(seed=42)

model.add(ae.Conv2d(3, 32, (3, 3), (1, 1), padding="same"))
model.add(ae.BatchNorm(epsilon=1e-5, momentum=0.9))
model.add(ae.ReLU())
model.add(ae.Conv2d(32, 32, (3, 3), (1, 1), padding="same"))
model.add(ae.BatchNorm(epsilon=1e-5, momentum=0.9))
model.add(ae.ReLU())
model.add(ae.MaxPool2d((2, 2), (2, 2), padding="valid"))
model.add(ae.SpatialDropout(rate=0.1, seed=42))

model.add(ae.GlobalAvgPool())
model.add(ae.Dense(32, 10))

# --- Configure, move to GPU, train ---
model.configure(
    loss=ae.SoftmaxCategoricalCrossEntropy(label_smoothing=0.05),
    optimizer=ae.AdamW(lr=0.001, decay=1e-4, weight_decay=0.01),
    accuracy=ae.CategoricalAccuracy(),
    preprocessor=feature_pipeline
)

model.to("cupy")
model.set_precision(compute_dtype="float16")
model.finalize(input_shape=X_train.shape[1:], fuse_kernels=True)
model.train(
    X=X_train,
    y=y_train,
    epochs=4,
    batch_size=256,
    shuffle=True,
    validation_data=(X_test, y_test)
)

model.save(filepath="saved_models/cifar10_cnn.aether")

Key Features

Layers

Conv2d, Dense, BatchNorm, MaxPool2d, AvgPool2d, GlobalAvgPool, Dropout, SpatialDropout, ReLU, LeakyReLU, SoftMax, with He initialization and Philox counter-based PRNG for dropout masks.

Kernels

Hand-written CUDA and HIP kernels for every layer and component except Dense, with kernels JIT-compiled once per operation/shape and memoized after.

Kernel Fusion

Automatic graph-level operator fusion on CuPy, fusing BatchNorm and ReLU/LeakyReLU on both forward and backward passes.

Training

AdamW with decoupled weight decay, automated mixed precision, label smoothing, batch shuffling, and in-training validation.

Preprocessing

Compose pipelines with Rescale, StandardScaler, and ToTensor.

Serialization

.aether bundles a JSON architecture manifest with safetensors weights (no pickle). The loader treats files as untrusted: only Aether's own classes resolve, and weights that don't match the manifest are rejected.