NNx

Lightweight PyTorch training, evaluation, and visualization with first-class graph neural network support.

Transparent orchestration for durable experiments, with your models and step logic left in your hands.

NNx is a lightweight PyTorch toolkit for repeatable training, evaluation and visualization. It provides frozen configuration objects, a supervised train, evaluate and predict loop, callbacks, metrics, and content-addressed checkpoints that support resume. Your models and per-step logic stay replaceable, so one loop serves standard, graph and transformer networks and other training methods. NNx is for researchers and engineers who want transparent PyTorch code without a larger training platform.

1. Why NNx¶

NNx owns the code that repeats across training projects: the loop, checkpoint handling and metric reporting. You write the model and, when you need to, the per-step logic.

1.1. Capabilities¶

Training loop — NNModel train, evaluate and predict, with callbacks, schedulers, FP32/FP16/BF16 precision, gradient accumulation and clipping, declared metrics, and a multi-optimizer Trainer.
Runs and resume — content-addressed run directories, checkpoint tags, verified warm resume, experiment plans and provenance records.
Tasks — classification, regression, multilabel, ranking, graph classification and link prediction.
Training methods — custom step functions and objectives, distillation, contrastive learning, Mixup and CutMix, Mixture of Experts, diffusion, I-JEPA and DPO.
Fine-tuning and transforms — layer freezing, LoRA, DoRA, IA3, prefix and prompt tuning, quantization, pruning and model surgery.
Networks and export — feed-forward, convolutional, graph, transformer and ViT networks; a small language-model path; ONNX, Hub, safetensors and experimental GGUF export.
Evaluation and decisions — reproducible splits, calibration, abstention, multi-seed comparison and typed decisions.

The Capability reference describes each item with its entry points. Concepts explains how the parts fit together.

2. Where to next¶

2.1. Get running¶

Install — a PyPI release or the current source, and the optional extras.
Quickstart — a runnable example that trains a model in seconds.

2.2. Understand the design¶

Concepts — runs and their on-disk layout, the enum factories, and how the specialization subpackages compose.

2.3. Deep-dive guides¶

Language modeling — train a tiny TransformerNN end-to-end (CPU-friendly).
Direct Preference Optimization — fine-tune an LM on (prompt, chosen, rejected) preference pairs.
I-JEPA — masked-patch latent-prediction self-supervised pretraining.
Model surgery — function-preserving Net2Net + low-rank + drop primitives.
Embeddings + FAISS — contrastive training + RAG-ready export.
HuggingFace Hub — safetensors + Hub publish/load.
Experimental GGUF export — container inspection, quantization notes, and runtime limits.
Typed decisions — provider-neutral Choice / Boolean / Score questions, decision jobs and the benchmark.
Comparison vs Lightning / HF / fastai / Composer — scope-explicit decision matrix for picking the right PyTorch training toolkit.
Architecture — package relationships plus exact callback and durable-commit ordering.
External dependency contracts — frozen integrations, release publication, and intentionally gated external checks.

2.4. Look things up¶

Capability reference — the main capabilities with their entry points, grouped by task.
API Reference — auto-generated from docstrings.
Export conformance — what an ONNX export record proves, profile by profile.
Feature composition — which feature combinations are verified, refused or unverified.
Examples catalog — annotated index of the runnable scripts under examples/.
CONTRIBUTING — editable install, dev toolchain, PR workflow.
Security policy — supported versions and private reporting instructions.
CHANGELOG — release history and unreleased changes.

3. Status¶

Alpha. Before 1.0, bug fixes can change behavior (see the Changelog), but public APIs are renamed only when they cannot work as designed.