简单的Torch模型测试：ModuleNotFoundError: 没有名为'ultralytics.yolo'的模块

Question

我有一个模型，叫做 best.pt，我想运行它。这个模型需要一张图片作为输入，然后会输出一个字符串。

我已经安装了 ultralytics、torch 和 torchvision 这些库。

我的代码很简单：

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)

结果如下：

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'

我哪里做错了呢？

Answer 1

看起来这个权重文件 best.pt 是用 Ultralytics 这个工具训练出来的。你可以试试下面这个方法。

from ultralytics import YOLO

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