Can we add more epochs to already trained yolo nas model and start training from where i left off?

Question

I have a YOLO NAS model for animal detection. I ran the model for 25 epochs and have got the best.pth weights. I need to add more epochs, to train it more from where i left off. I have read somewhere that YOLO V5 have such an option. Does YOLO NAS have the same option? If so how can i implement it in colab?

PS. It took me 10 - 15 hours to train for 25 epochs. So I am tight on time. I am not sure whether I am doing something wrong, but I am training using A100 GPU in Colab and its taking this much time. Please advice. I have 17.8 GB of data which has around 38790 images, so i guess it makes sense to take that much time?

I tried looking through the YOLO NAS documentation and google searched it, but couldn't get any concrete ideas.

Answer 1

Yes, you can do it, here is the important code to do that (assume you use Yolonas small version):

from super_gradients.training import models

dataset_params = {
    'classes': ['animal1', 'animal2', animal3']
}

model = models.get('yolo_nas_s',
                        num_classes=len(dataset_params['classes']),
                        checkpoint_path="/best.pth")

trainer.train(model=model, 
              training_params=train_params,
              train_loader=train_data, 
              valid_loader=val_data)