Optimizers¶

This tutorial presents the various optimizers available in timm. We look at how we could use each of them using the timm training script and also as standalone optimizers for custom PyTorch training scripts.

The various optimizers available in timm are:

1. SGD
2. Adam
3. AdamW
4. Nadam
5. Radam
6. AdamP
7. SGDP
8. Adadelta
9. Adafactor
10. ADAHESSIAN
11. RMSprop
12. NovoGrad

And some more from apex like:

1. FusedSGD
2. FusedAdam
3. FusedLAMB
4. FusedNovoGrad

which are GPU-only.

timm also supports lookahead optimizer.

As is the usual format for timm, the best way to create an optimizer using timm is to use the create_optimizer factory method. In this tutorial we will look at how to train each of these models using each of these optimizers using the timm training script first and also as standalone optimizers for custom training script.

Using timm training script¶

To train using any of the optimizers simply pass the optimizer name using the --opt flag to the training script.

python train.py ../imagenette-320/ --opt adam

Since Lookahead technique can be added to any of the optimizers, we can train our models using Lookahead in timm, simply update the optimizer name by adding a lookahead_ prefix. Example for adam, the training script looks like:

python train.py ../imagenette-320/ --opt lookahead_adam

And that's really it. This way we can train models on ImageNet or Imagenette using all the available optimizers in timm.

As standalone optimizers for custom training script¶

Many a time, we might just want to use the optimizers from timm for our own training scripts. The best way to create an optimizer using timm is to use the create_optimizer factory method.

NOTE: create_optimizer in timm accepts args as the first argument. This args parameter is from ArgumentParser so we might have to mock it to create optimizer for our custom training script. The example below shows how to do this.

The arguments for create_optimizer are shown below:

def create_optimizer(args, model, filter_bias_and_bn=True) -> Union[Optimizer, Lookahead]:
    """
    Here, `args` are the arguments parsed by `ArgumentParser` in `timm` training script. 
    If we want to create an optimizer using this function, we should make sure that `args` has the 
    following attributes set: 

    args: Arguments from `ArgumentParser`:
    - `opt`: Optimizer name
    - `weight_decay`: Weight decay if any 
    - `lr`: Learning rate 
    - `momentum`: Decay rate for momentum if passed and not 0 

    model: Model that we want to train 
    """

Let's see how to mock the args below:

Training NN using timm optimizers¶

In this section we are going to try and experiment some of the various available optimizers and use them in our own custom training script.

We will store the losses for each of the optimizers and in the end visualize the loss curves to compare performance on Imagenette dataset using a resnet-34 model that we again create using timm.

Data Dir¶

The directory structure of the data dir looks something like:

imagenette2-320
├── train
│   ├── n01440764
│   ├── n02102040
│   ├── n02979186
│   ├── n03000684
│   ├── n03028079
│   ├── n03394916
│   ├── n03417042
│   ├── n03425413
│   ├── n03445777
│   └── n03888257
└── val
    ├── n01440764
    ├── n02102040
    ├── n02979186
    ├── n03000684
    ├── n03028079
    ├── n03394916
    ├── n03417042
    ├── n03425413
    ├── n03445777
    └── n03888257

Train and Val dataset¶

Let's now create our train and validation datasets and dataloaders using timm. For more docs on datasets, refer here.

These are the class names that we have in Imagenette. We list them here for easy visualization below:

Let's now visualize some of the images and classes that our in our dataset.

It's great practice to visualize the images straight from the train_loader to check for any errors.

Train one epoch¶

In this section we will create our custom training loop.

The AverageMeter class below averages the loss for easy visualization. If we didn't take the moving average, then the loss curve would be very rocky and bumpy to visualize.

NOTE: As an experiment, feel free to remove the loss_avg = AverageMeter() in train_one_epoch function below and try visualizing the loss curve to see the difference.

The function below defines our custom training training loop. Essentially, we take the inputs and targets from the the train_loader. Get the predictions by passing the inputs through the model. Calculate the loss function, perform backpropogation using PyTorch to calculate the gradients. Finally, we use the optimizer to take step to update the parameters and zero out the gradients.

Also, note that we store the moving average of the losses for each of the mini batch losses.append(loss_avg.avg) in a list called losses. Finally, we return a dictionary with the Optimizer name and the list losses.

Note that this train_one_epoch function accepts args. These are the mocked args that we have looked at before. This args parameter get's passed to optim_factory.create_optimizer to create the Optimizer.

Let's now pass in the various Optimizers. The training loop that we have created should take care of instantiating the Optimizer using the create_optimizer function.

We have set the learning rate to be 1e-4, weight decay and momentum both to be 0.

We also pass in lookahead_adam to showcase training using the Lookahead class in timm.

Comparing the performance of different optimizers¶

Finally, let's visualize the results to compare the performance. All the losses alongside the Optimizer passed were stored in losses_dict.

We can see that Adam and AdamP perform the best out of the available optimizers on Imagenette for the 1 epoch that we have trained our model for. After this, please feel free to run your own experiments! :)