CODEWITHSUNDEEP
pythonpytorchyolov8
Kalleni
Kalleni

Reputation: 79

How to save a YOLOv8 model after some training on a custom dataset to continue the training later?

I'm training YOLOv8 in Colab on a custom dataset. How can I save the model after some epochs and continue the training later. I did the first epoch like this:

import torch

model = YOLO("yolov8x.pt")
model.train(data="/image_datasets/Website_Screenshots.v1-raw.yolov8/data.yaml", epochs=1)

While looking for the options it seems that with YOLOv5 it would be possible to save the model or the weights dict. I tried these but either the save or load doesn't seem to work in this case:

torch.save(model, 'yolov8_model.pt')
torch.save(model.state_dict(), 'yolov8x_model_state.pt')

Upvotes: 4

Views: 30483

Answers (3)

Omar
Omar

Reputation: 1

Well the problem is that runs/detect is not always available. For example if you exceed GPU limit the environment will stop and remove the GPU backend, after restarting you won't find runs directory when mounting to the drive. The problem is solved in yolov5 with save_dir parameter but for yolov8 the only solution that I found is dividing the training epochs so that usage limits won't be reached and I make a backup of runs directory in my drive.

As a result, I can make my predictions or continue training based on the best.pt or the last.pt from the backup directory.

Upvotes: 0

Hariharan
Hariharan

Reputation: 371

"I am currently working on a project using YOLOv8. After training on a custom dataset, the best weight is automatically stored in the runs/detect/train/weights directory as best.pt. When I retrain the model, I use the best.pt weight instead of yolov8x.pt to train the model."

Upvotes: 6

Rukon
Rukon

Reputation: 87

I think just replacing the "yolov8x.pt" to your trained model will do the work The trained model will be saved in the results/run folder in the working dir.

model = YOLO('yolov8x.yaml').load('yolov8x.pt') # build from YAML and transfer weights
model.train(data="/image_datasets/Website_Screenshots.v1-raw.yolov8/data.yaml", epochs=1)

check this -> https://docs.ultralytics.com/modes/train/

Upvotes: 1

Related Questions