Fine-tuning Custom Model

Instead of training a new image classification model from scratch, we can build on an existing model. This process is called fine-tuning a custom model. Fine-tuning is a method that applies transfer learning to repurpose a model for other tasks. It’s usually a lot cheaper and faster to fine-tune an existing model than to start over with a new one. For example, we can train an initial model for 10 classes with open-source datasets and then fine-tune it with our datasets. The performance of the final model is usually a lot better than it would be by training a new image classification model.

Fine-tune a custom model

The training script comes with the initial-checkpoint argument, which accepts a path string to the base model. It will initialize and load the weight from the checkpoint.

Run the following command to fine-tune a custom model:

python train.py /app/dataset2 --model resnet50 --num-classes 4 --initial-checkpoint /app/resnet50_best.pth.tar

Note: Some arguments, such as num-classes and img-size, must be the same as the initial checkpoint. It will cause an error if there’s a mismatch.

Resume full model and optimizer

The initial-checkpoint argument will load the model weights upon model creation. The training will start from epoch 0.

We utilize the resume argument to load the full model. It continues the training from where it left off. It also retains the optimizer state.

python train.py /app/dataset2 --model resnet50 --num-classes 4 --resume /app/resnet50_best.pth.tar

Note: We require additional memory to resume a full model with the optimizer state.

Ignore optimizer state

There’s an option that prevents the optimizer from resuming when a model resumes. Pass in the no-resume-opt flag as follows:

python train.py /app/dataset2 --model resnet50 --num-classes 4 --resume /app/resnet50_best.pth.tar --no-resume-opt

Now, the script ignores the optimizer state when loading the model.