← back to the tree

NNx

Lightweight PyTorch toolkit for training, fine-tuning, visualizing, and exporting modern neural nets.

one training loop26 example scriptsONNX, GGUF & HF export

pip install thekaveh-nnx

Glassy neural-network layers converging into a glowing NNx wordmark — data flowing through a training stack.
Neural layers converging on the wordmark. Digital painting, 2026 — NNx repository.

NNx is the training loop you’d otherwise rewrite in every notebook. It’s a lightweight PyTorch toolkit for training, evaluation, and visualization, with first-class graph neural networks, pulled out of a real notebook-based ML project once the boilerplate started repeating. It builds networks from frozen-dataclass configs, runs the train/eval/predict loop, manages checkpoints, and fires a documented callback lifecycle, so the notebook on top can stay about the model. Fine-tuning, diffusion, distillation, and PEFT bolt on through the same extension point when a task needs them, without anyone editing the core.

Highlights

  • Two training entry points — NNModel for supervised work and a multi-optimizer Trainer for GAN and actor-critic setups, over one loop with callbacks, schedulers, AMP, gradient clipping, and seeded reproducibility
  • Content-addressed checkpoints — NNRun writes to runs/<id>/ after every epoch and survives a KeyboardInterrupt, saving at six tags with optimizer-state sidecars for warm resume
  • Pluggable training paradigms — knowledge distillation, SimCLR, Mixup/CutMix, sparse Mixture-of-Experts, I-JEPA, and DPO, each entering the loop through one train_step_fn hook rather than a fork
  • Parameter-efficient fine-tuning — LoRA, DoRA, IA3, prefix, prompt, and adapters, each persisting only its trainable delta
  • Quantization and pruning — INT8 weight-only PTQ and 8da4w QAT via torchao, plus magnitude and 2:4 semi-structured pruning
  • Function-preserving model surgery — Net2Net widen and deepen, layer drop, and low-rank factorization, so a resized model resumes training without an accuracy cliff
  • Language modeling — a decoder-only Transformer with RMSNorm, RoPE, SwiGLU, and a KV-cache, plus greedy, top-k, and top-p sampling
  • Visualization and diagnostics — Plotly confusion matrices and training curves, weight histograms, activation maps, Captum-backed input attribution, per-layer gradient-flow checks, and a fastai-style learning-rate finder

Stack & Architecture

Python on PyTorch, with PyTorch Geometric behind the graph networks — GCN, GraphSAGE, and GAT on a shared base. Config objects are frozen, kw-only, slotted dataclasses that round-trip through state()/from_state(), and a callback bus fires lifecycle hooks to listeners like EarlyStopping and TensorBoardCallback. The heavier machinery — PEFT, quantization, diffusion, GGUF export — ships as opt-in extras, so a plain install stays small. Trained models export to ONNX, safetensors, GGUF, and the Hugging Face Hub, and a contrastive embeddings trainer turns (anchor, positive) text pairs into a domain-specific embedder exported as a FAISS index any RAG stack can consume.

Status

Actively developed and Apache-2.0 licensed, with a full documentation site that includes a generated API reference and a catalog of 26 runnable example scripts, ordered from the core loop out to DPO, pruning, and surgery.