NNx: A Lightweight PyTorch Neural Network Training Toolkit with Native Support for Graph Neural Networks

NNx is a lightweight PyTorch training framework distilled from experimental code. It provides a unified interface supporting feedforward neural networks and graph neural networks (GCN, GraphSAGE, GAT). Through dataclass configurations, it enables pluggable models, loss functions, optimizers, and schedulers, with built-in automatic checkpointing, visualization, and experiment reproducibility capabilities.

• Original Author / Maintainer: thekaveh
• Source Platform: GitHub
• Original Title: NNx
• Original Link: https://github.com/thekaveh/NNx
• Publication Date: 2026-05-25

---

During deep learning experiments, researchers often need to repeatedly write similar code for training loops, checkpoint saving, and result visualization. NNx is a lightweight toolkit distilled from such repetitive work. Initially serving as the underlying support for the thekaveh/ml project, it has now become an independent dedicated framework for neural network training and evaluation.

Unlike large frameworks like PyTorch Lightning or Hugging Face Transformers, NNx takes a "lean and focused" approach: it does not attempt to cover all possible deep learning scenarios, but instead focuses on the most common needs—training workflows for feedforward neural networks and graph neural networks—while maintaining code readability and modifiability.

NNx's design philosophy centers on several key abstractions, each corresponding to a core component in the deep learning workflow.

NNModel is the core orchestrator of the entire framework. It is responsible for building networks based on configuration parameters, executing training and evaluation, managing checkpoints, and supporting callback mechanisms (including early stopping). Its features include:

• Support for mixed-precision training
• Gradient clipping and gradient accumulation
• Seed fixing to ensure experiment reproducibility
• Warm-start training recovery capability

This centralized design allows users to avoid jumping between multiple modules; all training-related operations are completed through the unified interface of NNModel.

The framework has built-in multiple network implementations, all sharing a unified base class design:

• FeedFwdNN: A traditional feedforward network suitable for image and tabular data
• GraphConvNN: A graph neural network implementation based on Graph Convolutional Networks (GCN)
• GraphSageNN: An implementation of the GraphSAGE algorithm, suitable for large-scale graph data
• GraphAttNN: A Graph Attention Network (GAT) that supports multi-head attention mechanisms

All these networks inherit from GraphNNBase, with differences only in the specific selection of PyTorch Geometric layer constructors. This design greatly reduces the cognitive cost of switching network architectures.

NNx provides three main dataset encapsulations, each corresponding to different data types:

• NNDataset: Encapsulation based on torchvision's VisionDataset, handling traditional image data
• NNGraphDataset: Single-graph encapsulation based on PyG, using NeighborLoader for batch sampling
• NNTabularDataset: Converts pandas DataFrames into training/validation/test loaders

This layered design allows users to process image, tabular, and graph-structured data with almost the same code, without needing to rewrite data loading logic for each data type.

A key feature of NNx is its comprehensive dataclass configuration system. Each adjustable parameter has a corresponding frozen dataclass:

• NNParams: Network architecture parameters
• NNModelParams: Model orchestration parameters
• NNTrainParams: Training process parameters
• NNOptimParams: Optimizer configuration
• NNSchedulerParams: Learning rate scheduler configuration

These configuration objects support serialization and deserialization via .state() and .from_state() methods, ensuring that experiment configurations can be fully saved and restored. This design not only improves code readability but also provides a structured foundation for experiment management.