pyroml 2.0.0

Machine Learning framework allowing plug-and-play training for pytorch models

🔥 pyro

Lightweight Machine Learning framework allowing plug-and-play training for Pytorch models

• ⚡ Lightning inspired
• 💾 Support for wandb and checkpoints out-of-the-box
• 📊 Pretty logs, plots and support for metrics
• ✨ Fully type-safe
• 🪶 Lightweight and easy to use

Requirements

• Python 3.10 : 3.12

Installation

pip install pyroml

Locally

# Clone the repo
git clone https://github.com/peacefulotter/pyroml.git
cd pyroml

# Install dependencies
poetry config virtualenvs.in-project true # Optional, easier for vscode to find the venv folder
poetry install

Running tests

$ cd tests
$ python main.py # this will launch the training, follow the wandb link to access the plots
$ python pretrain.py # will load the last checkpoint and compute mse on a small part of the dataset, outputs True if model predicts correctly!