Simple Torch Model Test: ModuleNotFoundError: No module named 'ultralytics.yolo'

Question

I have a model, best.pt, that I'd like to run. It takes an image as input, and outputs a string.

I have ultralytics, torch and torchvision installed.

My code is simple:

import torch
from PIL import Image

# Load the pre-trained model
model = torch.load('best.pt')
# Load the input image
input_image = Image.open('input_image.jpg')
# Pass the image through the model
output = model(input_image)
# Print the output
print(output)

The result is as follows:

Traceback (most recent call last):
  File "/Users/fares/project/model/main.py", line 5, in <module>
    model = torch.load('best.pt')
            ^^^^^^^^^^^^^^^^^^^^^
  File "/Library/Frameworks/Python.framework/Versions/3.11/lib/python3.11/site-packages/torch/serialization.py", line 1026, in load
    return _load(opened_zipfile,
           ^^^^^^^^^^^^^^^^^^^^^
  File "/Library/Frameworks/Python.framework/Versions/3.11/lib/python3.11/site-packages/torch/serialization.py", line 1438, in _load
    result = unpickler.load()
             ^^^^^^^^^^^^^^^^
  File "/Library/Frameworks/Python.framework/Versions/3.11/lib/python3.11/site-packages/torch/serialization.py", line 1431, in find_class
    return super().find_class(mod_name, name)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
ModuleNotFoundError: No module named 'ultralytics.yolo'

What am I doing wrong?

Answer 1

It looks like the weights file best.pt was trained using the Ultralytics package. Give the following a try.

from ultralytics import YOLO

model = YOLO('best.pt')
input_image = Image.open('input_image.jpg')
output = model(input_image)
print(output)