Quickstart

Quickstart

An end-to-end CPU example you can paste into a Python REPL or script. Trains a tiny feed-forward classifier on random data so you can verify your install and understand the core flow — data → model → train → use — in one screen of code.

Minimal example

import torch
from torch.utils.data import DataLoader, TensorDataset

from nnx import (
    NNModel, NNParams, NNModelParams, NNTrainParams,
    NNOptimParams, NNSchedulerParams,
    Activations, Devices, Losses, Nets, Optims,
    EarlyStopping,
)

# 1. Data
X_train, y_train = torch.randn(256, 8), torch.randint(0, 3, (256,))
X_val,   y_val   = torch.randn(64,  8), torch.randint(0, 3, (64,))
train_loader = DataLoader(TensorDataset(X_train, y_train), batch_size=32, shuffle=True)
val_loader   = DataLoader(TensorDataset(X_val,   y_val),   batch_size=32)

# 2. Model
net_params   = NNParams(input_dim=8, output_dim=3, hidden_dims=[32, 16],
                        dropout_prob=0.1, activation=Activations.RELU)
model_params = NNModelParams(net=Nets.FEED_FWD, device=Devices.CPU,
                             loss=Losses.CROSS_ENTROPY)
model = NNModel(net_params=net_params, params=model_params)

# 3. Train
train_params = NNTrainParams(
    n_epochs=10,
    seed=42,                       # reproducibility
    train_loader=train_loader,
    val_loader=val_loader,
    optim=NNOptimParams(name=Optims.ADAM, max_lr=1e-2,
                        momentum=(0.9, 0.999), weight_decay=5e-5,
                        grad_clip_norm=1.0),
    scheduler=NNSchedulerParams(min_lr=1e-7, factor=0.5,
                                patience=3, cooldown=1, threshold=1e-3),
)
run = model.train(params=train_params, callbacks=[EarlyStopping(patience=5)])

# 4. Use it
print(f"trained {len(run.idps)} iterations; saved under runs/{run.id}/")
result = model.predict(X=X_val)
print(f"predicted {len(result.classes)} samples")

What just happened:

• NNParams describes the network shape (input width, output classes, hidden layers, dropout, activation).
• NNModelParams wires it to a loss function (Losses.CROSS_ENTROPY) and a device.
• NNModel constructs the FeedFwdNN net from those params; model.train() runs the loop.
• NNTrainParams carries the optimizer (NNOptimParams), scheduler (NNSchedulerParams), data loaders, epoch count, and seed.
• EarlyStopping stops training if validation error doesn't improve for 5 epochs.
• run is an NNRun with the iteration history; runs/<run.id>/ on disk holds checkpoints and idps.csv.
• model.predict(X=...) returns a PredictResult with .logits and .classes. The model is temporarily set to eval() mode and restored to train() afterward — calling predict mid-training is safe.

Common variations

GPU / Apple Silicon

from nnx import Devices
NNModelParams(net=Nets.FEED_FWD, device=Devices.get(), loss=Losses.CROSS_ENTROPY)
# Devices.get() picks MPS (Apple) > CUDA > CPU automatically.

Mixed precision (CUDA only)

NNModelParams(..., mixed_precision=True)   # silently no-op on CPU and MPS

Reproducibility

Pin every random-number generator before constructing data, model, or loaders:

from nnx import set_seed, dataloader_worker_init_fn

set_seed(42)                                       # pins torch / numpy / Python / cuDNN
loader = DataLoader(..., worker_init_fn=dataloader_worker_init_fn)
NNTrainParams(seed=42, ...)                        # also pins inside train() at entry

set_seed(42, strict=True) additionally enables torch.use_deterministic_algorithms(True) for bit-for-bit identical results across runs on the same hardware (may raise on ops without a deterministic CUDA kernel). See Reproducibility-and-Diagnostics.

Warm-resume training

NNTrainParams accepts resume_from_run_id to continue a previous run from a checkpoint. Optimizer state (Adam momentum, SGD velocity) is preserved via a .opt.pt sidecar:

# Round 1
run = model.train(params=NNTrainParams(n_epochs=10, ...))

# Round 2 — pick up from where the last run's LAST checkpoint left off.
NNModel(net_params=..., params=...).train(params=NNTrainParams(
    n_epochs=10,
    resume_from_run_id=run.id,
    resume_from_checkpoint="last",   # or "best" / "first" / "q1" / "q2" / "q3"
    ...
))

See Persistence-Runs-and-Checkpoints for the full checkpoint layout.

Loading a finished run

from nnx import NNRun, NNCheckpoint, Checkpoints, NNModel

run  = NNRun.load(id="<md5>")                          # rehydrate idps + params from disk
ckpt = NNCheckpoint.load(run=run.id, type=Checkpoints.BEST)
model = NNModel.from_checkpoint(checkpoint=ckpt)

Custom metrics

Pass a dict of callables to extra_metrics; each receives (y_true, y_pred) arrays and must return a scalar:

from sklearn.metrics import roc_auc_score

NNTrainParams(
    ...,
    extra_metrics={
        "roc_auc": lambda y, y_hat: float(roc_auc_score(y, y_hat, multi_class="ovr")),
    },
)
# Every NNEvaluationDataPoint gets `.extra["roc_auc"]` populated.
# Survives NNRun.load — the dict is stored in idps.csv.

Switching networks

Change the Nets value in NNModelParams; NNModel constructs the right network from the same NNParams:

NNModelParams(net=Nets.GRAPH_CONV,  device=Devices.CPU, loss=Losses.CROSS_ENTROPY)
NNModelParams(net=Nets.GRAPH_SAGE,  device=Devices.CPU, loss=Losses.CROSS_ENTROPY)
NNModelParams(net=Nets.GRAPH_ATT,   device=Devices.CPU, loss=Losses.CROSS_ENTROPY)
# For GRAPH_ATT, also pass n_heads= in NNParams.
# Feed batches via NNGraphDataset (PyG NeighborLoader-backed).

See Networks for all available network types and their params.

Scheduler choices

The default scheduler is ReduceLROnPlateau. Use NNSchedulerParams.builder() to switch to one of the five variants without hand-picking which fields belong to which kind:

from nnx import NNSchedulerParams, Schedulers

# Direct kwarg: all five required fields (min_lr, factor, patience, cooldown, threshold) must be supplied
NNSchedulerParams(kind=Schedulers.COSINE_ANNEALING, T_max=100, min_lr=1e-7, factor=0.5, patience=10, cooldown=2, threshold=1e-3)

# Or fluent builder (prevents invalid kind/field combinations)
scheduler = NNSchedulerParams.builder().one_cycle(
    max_lr=1e-3, total_steps=10_000,
    min_lr=1e-7, factor=0.5, patience=10, cooldown=2, threshold=1e-3,
).build()

Available kinds: REDUCE_LR_ON_PLATEAU (default), STEP, COSINE_ANNEALING, ONE_CYCLE, LINEAR_WARMUP_DECAY. See Fluent-Builders for the full builder API.

TensorBoard logging

pip install "thekaveh-nnx[tensorboard]"

from nnx import TensorBoardCallback
model.train(params=..., callbacks=[TensorBoardCallback(log_dir="tb_logs")])

LR finder pre-flight

Run an exponential learning-rate sweep before a long training run. The sweep is non-destructive — model weights and training mode are snapshotted and restored on exit:

import torch.nn.functional as F
from nnx import lr_finder

result = lr_finder(
    model.net, train_loader,
    loss_fn=F.cross_entropy,
    start_lr=1e-7, end_lr=10.0, num_iter=100,
)
print(f"Suggested max_lr: {result.suggested_lr:.2e}")
result.figure.show()   # Plotly: loss vs log(LR) with the suggestion marked

# Plug into the real run:
NNTrainParams(..., optim=NNOptimParams(name=Optims.ADAM, max_lr=result.suggested_lr, ...))

Silencing the progress bar (CI / non-TTY)

NNX_TQDM_DISABLE=1 python your_train_script.py

Any value of 1 / true / yes (case-insensitive) disables the tqdm bar in both NNModel.train() and Trainer.train().

Beyond supervised classification

When your loss is not loss_fn(net(X), Y) — autoencoder reconstruction, VAE composite loss, link prediction with negative sampling, diffusion noise prediction — pass a custom train_step_fn to model.train(). The rest of the loop (scheduler, callbacks, checkpoint cadence, val loop) stays the same. See Training-Loop-and-Callbacks and Training-Paradigms for details and factory functions that cover the common cases.

Reference

Symbol	Type	Where it lives
NNModel	class	nnx
NNParams	frozen dataclass	nnx
NNModelParams	frozen dataclass	nnx
NNTrainParams	frozen dataclass	nnx
NNOptimParams	frozen dataclass	nnx
NNSchedulerParams	frozen dataclass	nnx
EarlyStopping	Callback	nnx
predict	method on NNModel	NNModel.predict(X=...)
resume_from_run_id	field on NNTrainParams	str — md5 run id
extra_metrics	field on NNTrainParams	dict[str, Callable]
set_seed	function	nnx

See also

• Installation · Core-Concepts-and-Architecture · Training-Loop-and-Callbacks · Persistence-Runs-and-Checkpoints · Reproducibility-and-Diagnostics · Fluent-Builders · API reference · examples/