Fine-tuning Custom Model
Learn to fine-tune a custom image classification model.
We'll cover the following
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
andimg-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.
Get hands-on with 1400+ tech skills courses.