torchkeras is a PyTorch training utility designed to simplify and enhance the model training process, inspired by the user-friendly interface of Keras. It aims to provide an elegant, easy-to-use, and customizable training experience for PyTorch users, from researchers to practitioners, by abstracting away the complexities of manual training loops.

How It Works

torchkeras offers a KerasModel class that encapsulates the training logic, mirroring Keras's compile and fit API. Users define their model, loss function, optimizer, and metrics, then call model.fit() with their data loaders. The library handles the training loop, including epoch and step management, metric calculation, and callbacks, with built-in support for progress bars, visualizations, and early stopping. For advanced customization, users can override the StepRunner or EpochRunner classes to inject custom logic into specific stages of the training process.

Quick Start & Requirements

• Install: pip install torchkeras
• Prerequisites: PyTorch, tqdm, accelerate, torchmetrics, matplotlib, tensorboard, wandb (for specific features). GPU and CUDA are recommended for performance.
• Setup: Minimal setup required, primarily involves defining your PyTorch model and data loaders.
• Links: Examples

Highlighted Details

• Keras-like API: Simplifies PyTorch training with compile and fit methods.
• Built-in Visualization: Offers beautiful logs and dynamic visualizations in Jupyter notebooks.
• Extensibility: Supports callbacks for TensorBoard and Weights & Biases, and allows custom StepRunner and EpochRunner implementations.
• Accelerate Integration: Seamlessly supports GPU, multi-GPU (DDP), and mixed-precision (FP16/BF16) training via the accelerate library.

Maintenance & Community

The project is actively maintained by lyhue1991. Community interaction is encouraged via GitHub Issues and a WeChat group (details in README).

Licensing & Compatibility

The project does not explicitly state a license in the provided README. Compatibility for commercial use or closed-source linking would require clarification on the licensing terms.

Limitations & Caveats

The README does not specify a license, which may pose a barrier for commercial adoption or integration into closed-source projects. While the core KerasModel is straightforward, advanced customization requires understanding and modifying the StepRunner or EpochRunner components.