Pytorch Template 🔧

A custom template for scalable PyTorch projects with flexible configuration management. Designed to be ready in minimal time while remaining maintainable and scalable. ♻️

⚠️ This template is still in beta and may undergo significant changes. Feedback and contributions are welcome!

Template Structure 📂

PytorchTemplate/
├── readme.md  👈 You are here
├── requirements.txt            
├── run.py 
├── configs/ 
│   ├── config.py             
│   ├── default.yaml
│   └── ...
├── dataloaders/ 
│   └── ...                      
├── experiments/ 
│   ├── exp_base.py              
│   └── ...
├── models/                      
│   ├── base.py                  
│   └── ...
└── utils/                       
    ├── logger.py
    ├── binary_metrics.py
    ├── reproducibility.py
    └── ...                 

• run.py: Main entry point that handles configuration loading and experiment execution.
• configs/: The certainly_not_over_engineered configuration system (read below for more details).
• experiments/: Modular experiment classes that define training/inference logic. For larger projects, you may want a file for each experiment, such as forecasting, classification, etc.
• models/: PyTorch model implementations with a common base class that inherits all methods from torch.nn.Module.
• dataloaders/: Data loading and preprocessing modules.
• utils/: Shared utilities for logging, metrics, plotting, and reproducibility. Wandb comes pre-integrated in the logger.

Installation 🧨

1. Clone the repository or create it from this template:

   git clone <your-repo-url>
   cd <your-repo-name>

2. Install dependencies:

   python3 -m venv .venv
   pip3 install -r requirements.txt

🚯 To remove all files regarding the toy problem below and start with a freshly minted template run the following command!

bash cleanup.sh

Training and Inference 🚀

The template comes with a simple experiment inspired by the MinGPT Demo. From a sequence of 6 numbers, the transformer model should learn to sort them in ascending order. Input: 0 0 2 1 0 1 -> Output: 0 0 0 1 1 2

To train the model, run:

python run.py --config default.yaml toySorter.yaml --task training

The models will be trained in a minute or two, and the accuracy should reach 100%. To infer from the model, run:

python run.py --config default.yaml toySorter.yaml --task inference

Configuration 🔧

Configuration is handled by the configs/config.py file, which provides:

• Default values for all parameters.
• YAML configuration file loading.
• Command-line argument overrides.

This means that the configuration priority is (from highest to lowest): (1) command-line arguments, (2) YAML configuration file values, and (3) default values in the Config class.

Notes 📝

The template is developed as my personal starting point for new PyTorch projects. There is a trade-off when developing this type of template: the goal is to write the maximum amount of reusable code to save time in the future while avoiding to add complexity and knowledge overhead. I hope this strikes the right balance. It should be easy to understand and extend, while also providing a solid foundation for scalable projects.

The template will be updated as I discover new features to add or encounter bugs. If you have any suggestions or issues, feel free to open an issue on the repository.

License 📜

MIT License