Reputation: 21
I am currently trying to train a face detection model and have been using YOLOv8-face as found here: https://github.com/akanametov/yolov8-face The model works really well on its own, however I cannot seem to train it on my own image set without it overwriting everything that the pretrained model already knows.
So when I run without training, works great, detects faces. However when then training on my created yaml dataset of my own images, it all of a sudden forgets everything the model already knew and only has learning from my new dataset. How can I retain all of the pretrained models knowledge while adding to that and training on my own images to make it stronger at detection. here is my current code:
from ultralytics import YOLO
from PIL import Image
import cv2
import numpy as np
import requests
from io import BytesIO
model = YOLO("yolov8n-face.pt") #loading yolov8 face model
dataset = 'path(removed for privacy, but it works)' #specify dataset path
folder_path = 'test5' #specify image folder path
img_path = 'test_images/3251_Out.jpg' #if wanting to only predict single image, #replace folder_path with img_path in predict
print('model loaded')
model.train(data = dataset, epochs = 3, pretrained = True, imgsz=960) #training model
model.val() #evaluating train data
print('val complete')
#add this to predict when predicting on our set imgsize = 2160
results = model.predict(source=folder_path, save=True, conf = 0.33, line_width = 1, imgsz=2160)
Upvotes: 1
Views: 2659
Reputation: 1
For transfer learning in yolo v8 you have freeze a few initial layers and then then train your model on top of your pre-trained one. Try this :
model.train(data = dataset, epochs = 3, pretrained = "path to your pre-trained model", freeze = 5, imgsz=960)
Upvotes: 0