This repo contains a code structure that can be used as a simple PyTorch project template. It defines classes for dataset creation, data loader creation (including training/validation split handling), model, training procedure definition (including basic tensorboard logging).
To use this repo for your projects, you can fork/download this repo and use it as a starting point. To goal is to avoid to rewriting common steps such as
Dataset
toDataLoader
conversion- Validation/training set split logic
- Training loop definition
- Model saving and loading procedure
- ...
every time you start a new project.
Of course there is no such thing as a one-fits-all code. Hence the idea of the PyTorch project template, is that you derive child classes from the proposed base classes to override some functions according to your specific needs. To show how an example of such "application-specific adaptation" of the project template. I have included some example classes that are derived from the base classes for the CIFAR10 classification task. Note that, these classes are only here to showcase the functionalities of the project template, not to actually achieve any good classification results. The model is the same example model as in the PyTorch classification tutorial When you use this project template, you should create your own "example" classes specific to your task, needs and wishes.
-
data_handling
folder:base.py
: This file contains the definition of two base classes derived from PyTorch Dataset and DataLoader classes.BaseDataset
is here a pure abstract class, all its members need to defined according to the specific shape of the data. Seeexample_CIFAR10Dataset.py
for an example.BaseDataLoader
is derived fromDataLoader
, it can be instantiated as it is. It handles validation/training split of the input dataset in which case it has two attributes: train_loader and val_loader. Seeexample_main.py
for a usage example in the CIFAR10 classification task.example_CIFAR10Dataset.py
: This file contains an example of an application specific redefinition of theBaseDataset
class. It is build to transform the data when the data is in form of one folder containing the images where theid
is the name of the file and labels are stored in a csv. It correspond to the format of the data downloadable from Kaggle CIFAR10utils.py
: Use this file to put any data handling helper function you may need.
-
models
folder:base.py
This file contains the base class to build models. It implements thesave()
andload()
methods common to most models. Theinit()
andforward()
function need to be overridden for your specific model. Seeexample_simple_net.py
for an dummy example.example_simple_net.py
: dummy example of extension of theBaseModel
class
-
trainer
folder:base_trainer.py
This module contains a base implementation for a training procedure wrapped in aBaseTrainer
class. It defines atraining_step()
method, avalidation_step()
method and atrain()
method which consists of the main training loop. Note that the base trainer integrates basic tensorboard support (training loss and validation loss are recorded).base_predictor.py
This file contains the base implementation of theBasePredictor
class. This class can be used to iterate through a given test dataset and save the corresponding predictions. NOTE: This class is intended for testing NOT for validation as it assumes there are not labels available for the test set.example_CIFAR10Predictor.py
This file contains an example of an application specific redefinition of theBasePredictor
class.
-
example_main.py
This file contains an very simple example main file to show how to use this PyTorch Project template to build, train, validate, save, load a model and save predictions to a.csv
file.
Note: the plan is to further update this template to e.g. handle loading of model parameters via yaml
config files etc... whenever i'll find the time.