Загрузка YOLOv5 из PyTorch Hub
В этом руководстве объясняется, как загрузить YOLOv5 из PyTorch Hub — репозитория предобученных моделей для PyTorch.
Перед началом#
Установи requirements.txt в среде Python>=3.8.0, включая PyTorch>=1.8. Модели и наборы данных загружаются автоматически из последнего релиза YOLOv5.
pip install -r https://raw.githubusercontent.com/ultralytics/yolov5/master/requirements.txtКлонировать ultralytics/yolov5 не требуется — PyTorch Hub автоматически получает код.
Загрузка YOLOv5 через PyTorch Hub#
Простой пример#
В этом примере предобученная модель YOLOv5s загружается из PyTorch Hub как model, после чего ей передаётся изображение для инференса. 'yolov5s' — одна из небольших и быстрых моделей среди предобученных чекпойнтов YOLOv5.
import torch
# Model
model = torch.hub.load("ultralytics/yolov5", "yolov5s")
# Image
im = "https://ultralytics.com/images/zidane.jpg"
# Inference
results = model(im)
results.pandas().xyxy[0]
# xmin ymin xmax ymax уверенность класс название
# 0 749.50 43.50 1148.0 704.5 0.874023 0 person
# 1 433.50 433.50 517.5 714.5 0.687988 27 tie
# 2 114.75 195.75 1095.0 708.0 0.624512 0 person
# 3 986.00 304.00 1028.0 420.0 0.286865 27 tieПодробный пример#
В этом примере показан пакетный инференс с источниками изображений PIL и OpenCV — двумя типами входных данных, которые принимает метод forward YOLOv5 AutoShape(). Результаты results можно вывести в консоль, сохранить в runs/hub, показать на экране в поддерживаемых окружениях, а также вернуть в виде тензоров или датафреймов pandas.
import cv2
import torch
from PIL import Image
# Model
model = torch.hub.load("ultralytics/yolov5", "yolov5s")
# Images
for f in "zidane.jpg", "bus.jpg":
torch.hub.download_url_to_file(f"https://ultralytics.com/images/{f}", f) # download 2 images
im1 = Image.open("zidane.jpg") # PIL image
im2 = cv2.imread("bus.jpg")[..., ::-1] # OpenCV image (BGR to RGB)
# Inference
results = model([im1, im2], size=640) # batch of images
# Results
results.print()
results.save() # or .show()
results.xyxy[0] # im1 predictions (tensor)
results.pandas().xyxy[0] # im1 predictions (pandas)
# xmin ymin xmax ymax confidence class name
# 0 749.50 43.50 1148.0 704.5 0.874023 0 person
# 1 433.50 433.50 517.5 714.5 0.687988 27 tie
# 2 114.75 195.75 1095.0 708.0 0.624512 0 person
# 3 986.00 304.00 1028.0 420.0 0.286865 27 tie
Параметры инференса#
Модели YOLOv5 имеют различные атрибуты инференса, например порог уверенности и порог IoU, которые можно задать следующим образом:
model.conf = 0.25 # порог уверенности NMS
model.iou = 0.45 # порог IoU для NMS
model.agnostic = False # NMS без учёта классов
model.multi_label = False # несколько меток на рамку в NMS
model.classes = None # (необязательный список) фильтрация по классу, например = [0, 15, 16] для людей, кошек и собак в COCO
model.max_det = 1000 # максимальное количество обнаружений на изображение
model.amp = False # инференс с автоматической смешанной точностью (AMP)
results = model(im, size=320) # пользовательский размер для инференсаУстройство#
После создания модели можно перенести на любое устройство:
model.cpu() # CPU
model.cuda() # GPU
model.to(device) # например, device=torch.device(0)Модели также можно сразу создавать на любом device:
model = torch.hub.load("ultralytics/yolov5", "yolov5s", device="cpu") # загрузка на CPUПеред инференсом входные изображения автоматически переносятся на устройство, на котором находится модель.
Отключение вывода#
Модели можно загружать без вывода сообщений с помощью _verbose=False:
model = torch.hub.load("ultralytics/yolov5", "yolov5s", _verbose=False) # загрузка без вывода сообщенийВходные каналы#
Чтобы загрузить предварительно обученную модель YOLOv5s с 4 входными каналами вместо стандартных 3:
model = torch.hub.load("ultralytics/yolov5", "yolov5s", channels=4)В этом случае модель будет состоять из предобученных весов, кроме самого первого входного слоя, форма которого больше не совпадает с формой предобученного входного слоя. Веса входного слоя останутся случайно инициализированными.
Количество классов#
Чтобы загрузить предварительно обученную модель YOLOv5s с 10 выходными классами вместо стандартных 80:
model = torch.hub.load("ultralytics/yolov5", "yolov5s", classes=10)В этом случае модель будет состоять из предобученных весов, кроме выходных слоёв, форма которых больше не совпадает с формой предобученных выходных слоёв. Веса выходных слоёв останутся случайно инициализированными.
Принудительная перезагрузка#
Если при выполнении описанных выше шагов возникнут проблемы, задай force_reload=True: это удалит существующий кэш и принудительно загрузит последнюю версию YOLOv5 из PyTorch Hub. Кэшированные копии находятся в ~/.cache/torch/hub; удаление этой папки даст тот же результат.
model = torch.hub.load("ultralytics/yolov5", "yolov5s", force_reload=True) # принудительная перезагрузкаИнференс по снимку экрана#
Чтобы выполнить инференс на экране компьютера:
import torch
from PIL import ImageGrab
# Model
model = torch.hub.load("ultralytics/yolov5", "yolov5s")
# Image
im = ImageGrab.grab() # сделать снимок экрана
# Inference
results = model(im)Инференс на нескольких GPU#
Модели YOLOv5 можно параллельно загружать на несколько GPU и выполнять инференс в отдельных потоках:
import threading
import torch
def run(model, im):
"""Performs inference on an image using a given model and saves the output; model must support `.save()` method."""
results = model(im)
results.save()
# Models
model0 = torch.hub.load("ultralytics/yolov5", "yolov5s", device=0)
model1 = torch.hub.load("ultralytics/yolov5", "yolov5s", device=1)
# Inference
threading.Thread(target=run, args=[model0, "https://ultralytics.com/images/zidane.jpg"], daemon=True).start()
threading.Thread(target=run, args=[model1, "https://ultralytics.com/images/bus.jpg"], daemon=True).start()Обучение#
Чтобы загрузить модель YOLOv5 для обучения, а не для инференса, задай autoshape=False. Чтобы загрузить модель со случайно инициализированными весами (для обучения с нуля), используй pretrained=False. В этом случае нужно предоставить собственный скрипт обучения или обучать модель с помощью train.py, как показано в руководстве по обучению на пользовательских данных для YOLOv5.
import torch
model = torch.hub.load("ultralytics/yolov5", "yolov5s", autoshape=False) # загрузить предварительно обученную модель
model = torch.hub.load("ultralytics/yolov5", "yolov5s", autoshape=False, pretrained=False) # загрузить модель с нуляРезультаты в Base64#
Результаты в кодировке Base64 полезны для API-сервисов, как показано в примере Flask REST API для YOLOv5.
import base64
from io import BytesIO
from PIL import Image
results = model(im) # inference
results.ims # массив исходных изображений (в виде массива np), переданных модели для инференса
results.render() # обновляет results.ims рамками и метками
for im in results.ims:
buffered = BytesIO()
im_base64 = Image.fromarray(im)
im_base64.save(buffered, format="JPEG")
print(base64.b64encode(buffered.getvalue()).decode("utf-8")) # изображение с результатами в кодировке base64Обрезанные результаты#
Результаты можно вернуть и сохранить в виде обрезанных областей обнаружения:
results = model(im) # inference
crops = results.crop(save=True) # словарь обрезанных областей обнаруженияРезультаты в Pandas#
Результаты можно вернуть в виде DataFrame Pandas:
results = model(im) # inference
results.pandas().xyxy[0] # DataFrame PandasВывод pandas (нажми, чтобы развернуть)
print(results.pandas().xyxy[0])
# xmin ymin xmax ymax уверенность класс название
# 0 749.50 43.50 1148.0 704.5 0.874023 0 person
# 1 433.50 433.50 517.5 714.5 0.687988 27 tie
# 2 114.75 195.75 1095.0 708.0 0.624512 0 person
# 3 986.00 304.00 1028.0 420.0 0.286865 27 tieСортировка результатов#
Результаты можно сортировать по столбцам, например, чтобы упорядочить обнаруженные цифры на номерном знаке слева направо (по оси x):
results = model(im) # inference
results.pandas().xyxy[0].sort_values("xmin") # отсортировано слева направоРезультаты в JSON#
Результаты можно вернуть в формате JSON, предварительно преобразовав их в датафреймы .pandas() с помощью метода .to_json(). Формат JSON можно изменить с помощью аргумента orient. Подробности см. в документации pandas .to_json().
results = model(im) # inference
results.pandas().xyxy[0].to_json(orient="records") # предсказания JSON для img1Вывод JSON (нажми, чтобы развернуть)
[
{
"xmin": 749.5,
"ymin": 43.5,
"xmax": 1148.0,
"ymax": 704.5,
"confidence": 0.8740234375,
"class": 0,
"name": "person"
},
{
"xmin": 433.5,
"ymin": 433.5,
"xmax": 517.5,
"ymax": 714.5,
"confidence": 0.6879882812,
"class": 27,
"name": "tie"
},
{
"xmin": 115.25,
"ymin": 195.75,
"xmax": 1096.0,
"ymax": 708.0,
"confidence": 0.6254882812,
"class": 0,
"name": "person"
},
{
"xmin": 986.0,
"ymin": 304.0,
"xmax": 1028.0,
"ymax": 420.0,
"confidence": 0.2873535156,
"class": 27,
"name": "tie"
}
]Пользовательские модели#
В этом примере загружается пользовательская модель YOLOv5s, обученная на 20 классах VOC, 'best.pt', с помощью PyTorch Hub.
import torch
model = torch.hub.load("ultralytics/yolov5", "custom", path="path/to/best.pt") # локальная модель
model = torch.hub.load("path/to/yolov5", "custom", path="path/to/best.pt", source="local") # локальный репозиторийМодели TensorRT, ONNX и OpenVINO#
PyTorch Hub поддерживает инференс в большинстве форматов экспорта YOLOv5, включая модели, обученные на пользовательских данных.
- TensorRT может работать в 2–5 раз быстрее PyTorch по результатам тестов на GPU.
- ONNX и OpenVINO могут работать в 2–3 раза быстрее PyTorch по результатам тестов на CPU.
import torch
model = torch.hub.load("ultralytics/yolov5", "custom", path="yolov5s.pt") # PyTorch
model = torch.hub.load("ultralytics/yolov5", "custom", path="yolov5s.torchscript") # TorchScript
model = torch.hub.load("ultralytics/yolov5", "custom", path="yolov5s.onnx") # ONNX
model = torch.hub.load("ultralytics/yolov5", "custom", path="yolov5s_openvino_model/") # OpenVINO
model = torch.hub.load("ultralytics/yolov5", "custom", path="yolov5s.engine") # TensorRT
model = torch.hub.load("ultralytics/yolov5", "custom", path="yolov5s.mlpackage") # CoreML (только для macOS)
model = torch.hub.load("ultralytics/yolov5", "custom", path="yolov5s.tflite") # TFLite
model = torch.hub.load("ultralytics/yolov5", "custom", path="yolov5s_paddle_model/") # PaddlePaddleПоддерживаемые среды#
Ultralytics предлагает широкий выбор готовых окружений с предварительно установленными основными зависимостями, такими как CUDA, cuDNN, Python и PyTorch, чтобы ты мог сразу приступить к проектам.
- Бесплатные ноутбуки с GPU:
- Google Cloud: краткое руководство по GCP
- Amazon: краткое руководство по AWS
- Azure: краткое руководство по AzureML
- Docker: краткое руководство по Docker
Статус проекта#
Этот значок означает, что все тесты YOLOv5 GitHub Actions для непрерывной интеграции (CI) проходят успешно. Эти тесты CI тщательно проверяют функциональность и производительность YOLOv5 по ряду ключевых направлений: обучение, валидация, инференс, экспорт и сравнительное тестирование. Они обеспечивают стабильную и надёжную работу в macOS, Windows и Ubuntu. Тестирование проводится каждые 24 часа и после каждого нового коммита.
Часто задаваемые вопросы#
Нет. При первом использовании
torch.hub.load("ultralytics/yolov5", ...)загружает код в~/.cache/torch/hubи скачивает веса. Установить нужно только пакеты изrequirements.txt.Используй точку входа
custom, указав путь к контрольной точке:torch.hub.load("ultralytics/yolov5", "custom", path="runs/train/exp/weights/best.pt"). Добавьsource="local"и путь к локальному репозиторию, чтобы загрузить модель без доступа к сети.Да. Передай экспортированный файл тому же загрузчику
custom, и YOLOv5 выберет подходящий бэкенд по расширению файла, напримерyolov5s.onnx,yolov5s.engineилиyolov5s_openvino_model/.Передай
force_reload=True, чтобы удалить закэшированный репозиторий и загрузить последнюю версию кода, или удали~/.cache/torch/hubвручную. Так ты также получишь новые выпуски YOLOv5.results.xyxy[0]возвращает тензор с[x1, y1, x2, y2, confidence, class]строками,results.pandas().xyxy[0]возвращает DataFrame с названиями классов, а.to_json(orient="records")для этого DataFrame формирует JSON.
Комментарии
Задай вопрос, поделись результатами или помоги кому-нибудь.