Making predictions with custom model

Support
1

Would like to check if I understood correctly about the need to load a new custom model before making predictions…

Following the video tutorial here Accelerating YOLO11 Projects with Google Colab I see that, after running model.train() we need to load the new custom model before we can start using it to make predictions, for example model_custom = YOLO("/path/to/custom/model.pt") to load the custom model, then you can do model_custom.predict().

However, the sample code at https://github.com/ultralytics/ultralytics suggests you can directly do:

# Load a pretrained YOLO11n model
model = YOLO("yolo11n.pt")

# Train the model on the COCO8 dataset for 100 epochs
train_results = model.train(
    data="coco8.yaml",  # Path to dataset configuration file
    epochs=100,  # Number of training epochs
    imgsz=640,  # Image size for training
    device="cpu",  # Device to run on (e.g., 'cpu', 0, [0,1,2,3])
)

# Evaluate the model's performance on the validation set
metrics = model.val()

# Perform object detection on an image
results = model("path/to/image.jpg")  # Predict on an image
results[0].show()  # Display results

In this sample code I think the predictions will be using the pretrained model yolo11n.pt and not the new custom model that was just trained. Did I understand that correctly?

2

Both would work. Loading the new model is recommended to clear the model object from training data. But it’s not required.

3

That’s a bit different to what I’m seeing. I’m seeing prediction classes that are different to those returned by model.names. Is that expected?

If I call model.train() then model.predict() the prediction classes are those of the original (pretrained) model. I expected that the prediction classes would be from the classes defined in the training data. If I check model.names I see the classes from the training dataset, but the prediction classes are from the original model’s dataset.

So, I know we can just run one more line of code and load the newly-trained model, but I find it a bit confusing when the model object’s class predictions are not taken from model.names.

This is the code I’m running in Google Colab:

!pip install ultralytics
!pip install roboflow

from roboflow import Roboflow
rf = Roboflow(api_key=[api_key_goes_here])
project = rf.workspace("conveyor-550m0").project("conveyor-hhrzw")
version = project.version(3)
dataset = version.download("yolov11")

from ultralytics import YOLO
model = YOLO("yolo11n.pt")

model.names # Classes are from pretrained model {0: 'person', 1: 'bicycle', 2: 'car', 3: 'motorcycle', ...

results = model.predict("/content/conveyor-3/valid/images/box_0878_png.rf.873c452bc17b385b0380a9bc4316d38f.jpg")
results[0].show() # Prediction classes are "chair" and "suitcase"

train_results = model.train(data="/content/conveyor-3/data.yaml", epochs=10)

model.names # Classes are from training data {0: 'cardboard box', 1: 'conveyor', 2: 'kartonbox'}

results = model.predict("/content/conveyor-3/valid/images/box_0878_png.rf.873c452bc17b385b0380a9bc4316d38f.jpg")
results[0].show() # Prediction classes are still "chair" and "suitcase", so not from model.names

# Now load the new model we just trained
model_best = YOLO("/content/runs/detect/train/weights/best.pt")

results = model_best.predict("/content/conveyor-3/valid/images/box_0878_png.rf.873c452bc17b385b0380a9bc4316d38f.jpg")
results[0].show() # Prediction class is "cardboard box"
4

That’s because you ran prediction before training which creates a predictor using the old model that is still retained in the subsequent prediction because we don’t want to create predictor every time.

You can run model.predictor = None if you want to reset the predictor.

1 Like