Search code examples
pythonpytorchyolov8

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')
Solution

  • "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."