This repository contains a framework for training deep learning-based classification and regression models with Pytorch Lightning.
It comprises several architectures, regularization, augmentation and training techniques and aims to provide easy-to-use baselines for experimenting with a lot of different setups.
You can also integrate your own model and/or dataset and benefit from the features of this repository!
Results of experiments on CIFAR-10 comparing different architectures in different training settings are shown below.
Everything can be run via the Command Line Interface and with Hydra config files. Logging is accomplished by Weights&Biases.
Training uses mixed precision, torch.compile and torch.backends.cudnn.benchmark=True by default to increase training speed.
Detailed CIFAR results and used configurations can be seen in CIFAR Results. Best performance is achieved with a PyramidNet using RandAugment augmentation, Shakedrop and Mixup. It yields 0.986 Test Accuracy on CIFAR-10 and 0.875 Test Accuracy on CIFAR-100.

• How to run
  • Requirements
  • General Usage & Hydra config files
  • Available models and parameters
    • Models
    • Training Settings
    • Data Settings
    • Regularization Techniques
  • Logging
• Including custom pytorch models
• Including other datasets
• CIFAR Results

First install the requirements in a virtual environment by:

pip install -r requirements.txt

You might need to adapt the cuda versions for torch and torchvision. Find a torch installation guide for your system here.

Everything in this repository can be run by executing the main.py script with a corresponding configuaration. We use Hydra for config management. All config files are stored in the cli_configs directory:

.
├── cli_configs/
│   ├── train.yaml
│   ├── data/
│   │   └── cifar10.yaml
│   │   └── imagenet.yaml
│   │   └── poster_regression.yaml
│   │   └── ...
│   ├── env/
│   │   └── local.yaml
│   │   └── cluster.yaml
│   └── model/
│       └── convnext.yaml
│       └── resnet.yaml
│       └── efficientnetv2.yaml
│       └── ...      
└── ...

The train.yaml defines all default values. The other 3 subdirectories define config groups that can be used to overwrite specific values. The three config groups are data, environment and model. Inside the data config group each dataset has its own yaml file defining any dataset specific settings like batch size, augmentations, number of classes etc. The environment config group defines the paths to the data directory as well as the experiment directory where logs are saved and sets environment specific variables, e.g. the progress bar gets disabled in a cluster environment. The model config group contains configs for all available models.

You can start training by simply calling the main.py and add the config groups you want. Here is an example of training an EfficientNetV2 on CIFAR-10 in the local environment:

python main.py env=local model=efficientnetv2 data=cifar10

Each defined config variable can be overwritten either in the command line or by directly modifiying the yaml files for permanent changes. New yaml files can be added, e.g. a new env config defining paths for your specific environment. Final configurations depend on the order, i.e. a variable that is defined in several config files or in the command line will always get the last defined value.

For overwriting the default EfficientNetV2 model type S, you can directly change the value in the cli_configs/model/efficientnetv2.yaml file or simply use the command line:

python main.py env=local model=efficientnetv2 data=cifar10 model.type=M

It is possible to define default config groups by specifying them at the top of train.yaml. They are used when no other config groups are specified in the command line. By default they are set to:

defaults:
  - _self_          # all variables defined in train.yaml
  - env: local      # local environment
  - model: resnet   # ResNet18
  - data: cifar10   # training on CIFAR-10

If you want to include your own model please see Including custom models for instructions. See Including custom datasets for instructions on using your own datasets.

The following models and parameters are available out of the box:

Weights&Biases is used for logging all hyperparameters, metrics, terminal output, system metrics such as GPU utilization etc. You can create an account for free. However, if you don't want to create an account you can still use this repository. W&B will display a link in the beginning of training where you can find all your logs, but they will be deleted after 7 days if you don't have an account. We recommend to create an account.

The name of your dataset is automatically used as project name. You can adjust the project name by overwriting

trainer.logger.project=<PROJECT_NAME>

Moreover, additional plots depending on the task are saved. For classification a confusion matrix is automatically logged for the validation set. You can also log one for the train set or disable it completely by overwriting the default value val:

model.result_plot=all  # logs train and val plot
model.result_plot=disable  # does not log plots

In the case of a regression task a parity plot is logged showing ground truth and predicted regression targets instead of a confusion matrix.

For including your own models follow these steps:

1. In the ./models folder create a file <your-model>.py and import your model.

2. Create a class that inherits from the BaseModel and define the __init__ and forward functions. Initialize your model in the __init__ function. Here is an example for a model that expects the arguments num_classes, stochastic_depth_prob and some_new_argument:

   from my_custom_models import custom_model  # your model
   from base_model import BaseModel  # customized LightningModule
   
   class CustomModel(BaseModel):
       def __init__(self, **kwargs):
           super(CustomModel, self).__init__(**kwargs)
   
           self.model = custom_model(
               num_classes=kwargs["num_classes"],
               stochastic_depth_prob=kwargs["stochastic_depth"],
               some_new_argument=kwargs["some_new_argument"]
               ...  # any other arguments
           )
           
       def forward(self, x):
           return self.model(x)

   Note that **kwargs is passed to the __init__ and the super init. These are all arguments defined in the model config group (i.e. all variables starting with model.<variable>). You can use existing variables like num_classes or you can define new ones and simply add them to your config files.

3. Create a <your-model>.yaml file in the model config group. Add the path to your CustomModel created in step 2, default parameters you want to use with this model and if applicable default values for new arguments. Here is an example:

   _target_: models.custom_model.CustomModel
   stochastic_depth: 0.2
   some_new_argument: <some-value>

   Note that num_classes is not defined here since it is usually defined in the data config file.

That's it! Note that the techniques Stochastic Depth, Shakedrop, Squeeze & Excitation, Final Layer Dropout and Bottleneck require architectural changes and cannot be made available automatically for a custom model. All other functionalities of this repository are now also available for your model. You can now train your model e.g. with RandAugment, Mixup, Madgrad Optimizer and a Cosine Annealing LR scheduler with warmstart of 5 epochs all out of the box like this:

python main.py model=custom_model model.lr=0.0001 model.optimizer=Madgrad model.scheduler=CosineAnneal model.warmstart=5 model.mixup=True data.module._target_=augmentation.policies.cifar.AutoAugmentTransform

For including your own dataset follow these steps:

1. In the dataset directory create a new file that implements the torch dataset class for your data.
2. Additionally, create the DataModule for your dataset by writing a class that inherits from BaseDataModule. Write the init and setup functions for your dataset. The dataloaders are already defined by the BaseDataModule. An example could look like this:

   from .base_datamodule import BaseDataModule
   
   class CustomDataModule(BaseDataModule):
     def __init__(self, **params):
         super(CustomDataModule, self).__init__(**params)
   
     def setup(self, stage: str):
         self.train_dataset = YourCustomPytorchDataset(
             data_path=self.data_path,
             split="train",
             transform=self.train_transforms,
         )
         self.val_dataset = YourCustomPytorchDataset(
             data_path=self.data_path,
             split="val",
             transform=self.test_transforms,
         )

   Note that the __init__ function takes **params and passes them to the super init. By doing so the attributes self.data_path, self.train_transforms and self.test_transforms are already set automatically and can be used in the setup function. The self.data_path is a joined path consisting of the configs data.module.data_root_dir and data.module.name. Custom transforms can be added in ./augmentation/policies/<your-data>.py. They need to inherit from the BaseTransform class. See the existing transforms for examples!
3. Add a <your-data>.yaml file to the data config group, defining some data-specific variables. For CIFAR-10 it looks like this:

   # @package _global_
   data:
     module:
       _target_: datasets.cifar.CIFAR10DataModule
       name: CIFAR10
       batch_size: 128
       train_transforms: 
         _target_: augmentation.policies.cifar.RandAugmentTransform
         cutout_size: 8
       test_transforms: 
         _target_: augmentation.policies.cifar.TestTransform
   
     num_classes: 10
   
   model:
     task: 'Classification'
     cifar_size: True

   The data.module._target_ defines the path to your DataModule. Note that the first line of the file needs to be # @package _global_ in order for Hydra to read the config properly.

That's it! You can now use all training pipelines, regularization techniques and models with your dataset.

The following experiments aim to show general tendencies for different training settings and techniques. In order to be consistent with the literature we evaluate all our experiments on the respective test sets. We are aware of problems resulting from this experimental setup and would like to emphasize that test set optimization should never be done when developing new methods or making claims about the generality of the results. A proper experimental setup should always include a validation set which is used for hyperparameter optimization. Thus, all of the following experiments are merely intended to provide general tendencies without any guarantee that they will generalize to unseen data.

The following table shows the exact configurations that were used to achieve the stated test accuracies on CIFAR-10 and CIFAR-100. Note that not all possible combinations of parameters were tested. Therefore, one shouldn't draw any conclusions about single parameters. The table only aims to enable reproducibility. Please see the experiments in the following sections for an evaluation of individual techniques.

The choice of an optimizer, corresponding learning rate and learning rate scheduler all have a large impact on convergence speed and model performance. Here different setups were compared for the ResNet34 model. All versions were trained for 200 epochs, using a batch size of 128, AutoAugment augmentation and a weight decay of 4e-5. No other regularization techniques or additional settings were used. SGD with momentum of 0.9 and no nesterov was compared to Adam (no amsgrad), AdamW (no amsgrad), RMSProp and the recently introduced Madgrad. In addition, Sharpness Aware Minimization (SAM) was used with SGD and Madgrad as base optimizers. The considered learning rate schedulers were Step (multiply lr with 0.1 evey 50 epochs), MultiStep (multiply lr with 0.1 after 1/2 and 3/4 of epochs) and Cosine Annealing without warmstart. Every combination of LR, optimizer and scheduler that is not displayed yielded a score lower than 0.92 and is not shown to give a more clear view on the good performing setups.
For reproducing e.g. the best SGD - SAM run which yielded more than 0.97 accuracy you can run:

python main.py model=resnet model._target_=models.dynamic_resnet.ResNet34 model.lr=0.1 model.optimizer=SGD model.sam=True model.scheduler=CosineAnneal data.module._target_=augmentation.policies.cifar.AutoAugmentTransform trainer.max_epochs=200

It can be seen that optimizer and LR are highly dependent on each other. While SGD performs better with higher LRs like 0.1, the other ones yield better results with much lower LRs. Madgrad yields best results with the smallest LR tested (0.0001). Moreover, Sharpness Aware Minimization (SAM) always outperformed the respective base optimizer alone. Regardless of the optimizer the best scheduler was usually Cosine Annealing.

Baseline augmentation on CIFAR involves Random Cropping, Random Horizontal Flips (with a 0.5 probability) and Normalization. In addition, Cutout (with a size of 16 on CIFAR-10 and 8 on CIFAR-100) is added in baseline_cutout. It erases a random patch of the image and is able to improve the robustness of the model by simulating occlusions. AutoAugment and RandAugment both include the baseline augmentation, then their respective policy and finally also utilize cutout. While AutoAugment is a set of learned augmentation policies that work well on CIFAR, RandAugment does not have to be learned first but applies a subset of transformations randomly.

Here all models were using either SGD or SGD with SAM as an optimizer, an initial LR of 0.1 and a Cosine Annealing scheduler. While the SGD models were trained for 400 epochs, the SAM models were trained for 200 epochs to ensure a fair comparison (in terms of total training time). The PreActResNets were all using a gradually increasing warmstart for 10 epochs before starting to decay the LR, because otherwise the LR of 0.1 would have been too high for them in the beginning, so they would result in NaNs. The other models did not use a warmstart. If you e.g. want to reproduce the results of a PreActResNet50 with RandAugment and SGD - SAM optimizer, you can run the following:

python main.py model=preact_resnet model._target_=models.preact_resnet.PreActResNet50 model.lr=0.1 model.optimizer=SGD model.sam=True model.warmstart=10 trainer.max_epochs=200

The respective run with SGD alone would be:

python main.py model=preact_resnet model._target_=models.preact_resnet.PreActResNet50 model.lr=0.1 model.optimizer=SGD model.warmstart=10 trainer.max_epochs=400

Single runs are shown for each augmentation technique with their resulting test accuracy. Adding Cutout to the baseline considerably improves the performance of all models. AutoAugment and RandAugment yield very similar results but both consistently outperform the baselines with and without cutout. Moreover, when comparing the performance of the optimizers it can be seen that the influence of SAM varies between augmentations and model depth. In the case of baseline and cutout augmentation using SAM always improved the final test accuracy significantly, regardless of the model. However, when using heavier augmentation techniques like AuoAugment and RandAugment SAM only improved performance of deeper models. VGG16 as well as the (PreAct)ResNets18 and 34 all did not improve using SAM or even performed worse. On the other hand the PyramidNets and deeper ResNets still yielded better scores when utilizing SAM.

Here the previous runs using AutoAugment are shown with their respective time needed to complete a full training. In addition, SGD runs with only 200 epochs are depicted. All models were trained on a GeForce RTX 2080 Ti. As expected smaller models usually need less training time but yield lower accuracy. The WideResNet 28-10 however, trains faster than the ResNet152 and PreActResNet152 but still performs better. Training the plain SGD versions for 400 instead of 200 epochs always increases the performance significantly except for the ResNet152 and PreActResNet152. As seen in the previous section, SAM used with more complex models outperforms the 400 epoch SGD runs. Therefore, when using SAM training for only half of the epochs you would usually do can be a good trade-off between performance and training time.

There are several architectural tweaks such as e.g. stochastic depth, general regularization techniques (e.g. mixup) and other tricks like using a LR scheduler warmstart or using nesterov accelerated SGD. Find a short description, links to papers and an explanation how they were used for all techniques in Available models and parameters. In order to test their effects on the CIFAR datasets each of them was added individually to a baseline model. Since final test accuracies can vary for similar runs, all models were trained 5 times. The baseline models were trained for 200 epochs, using SGD with a LR of 0.1, a Cosine Annealing scheduler and AutoAugment augmentation. The mean test accuracies are depicted by their respective dashed vertical lines. The results of the runs including each respective technique are represented by boxplots to show the variance within runs. An example run for a WideResNet using mixup would be:

python main.py model=wide_resnet model.lr=0.1 model.optimizer=SGD model.scheduler=CosineAnneal data.module._target_=augmentation.policies.cifar.AutoAugmentTransform trainer.max_epochs=200 model.mixup=True

It can be seen that variance in runs is higher for some techniques such as squeeze & excitation while it is rather stable for other techniques such as nesterov or mixup. Moreover, different models work differently well with some techniques like final layer dropout which mostly increases performance of the WideResNet while it harms performance of the other two models. For all three models not decaying batchnorm and bias decreased test accuracy while using a warmstart, mixup and nesterov all increased performance.

This Repository is developed and maintained by the Applied Computer Vision Lab (ACVL) of Helmholtz Imaging.