Flame is a lightweight and efficient library for developing neural networks, built using PyTorch Ignite. It simplifies the training, evaluation, and experimentation processes by providing reusable templates and a highly configurable setup. With Flame, you can easily manage everything from logging and checkpoints to custom training loops, making it ideal for both beginners and experienced researchers looking to streamline their workflow.
Flame addresses two common needs in neural network development:
- Experiment Templates: Provides templates for training and evaluation with utilities like checkpoint saving, training resumption, logging, and more.
- Flexible Functionality: Easily customize training and testing by editing configuration files. Integrate with Ignite's metrics and handlers or use your own.
Details coming soon.
Set up a Python 3 environment and install Flame using pip
:
pip install git+https://github.com/tronglh241/flame.git
Flame includes two commands:
-
Initialize a new project:
flame init [-h] [-f] [directory]
directory
: Location to initialize the project. Defaults to the current directory.-f, --full
: Creates a full project template with an additionalrun.py
file.
-
Run training or testing:
flame run [-h] file
file
: Path to the config file.
Follow these steps to run a simple classification experiment using the MNIST dataset:
-
Initialize a project:
mkdir mnist-classification && cd mnist-classification flame init
Or, run directly:
flame init mnist-classification
The folder structure will look like this:
mnist-classification/ └── configs ├── test_ddp.yml ├── test.yml ├── train_ddp.yml └── train.yml
Use
-f
for an extrarun.py
file:mnist-classification/ ├── configs │ ├── test_ddp.yml │ ├── test.yml │ ├── train_ddp.yml │ └── train.yml └── run.py
-
Install additional dependencies:
pip install torchvision
-
Run the training:
flame run configs/train.yml
To monitor training, use Tensorboard:
tensorboard --logdir logs/
-
Evaluate the model: Update
checkpoint.loader.kwargs.path
inconfigs/test.yml
to point to the trained model, e.g.,checkpoints/best_model.pt
:checkpoint: loader: module: flame.handlers name: CheckpointLoader kwargs: path: "'checkpoints/best_model.pt'"
Run the evaluation:
flame run configs/test.yml
That's it! You've successfully trained and evaluated a model using Flame.