Загрузка YOLOv5 из PyTorch Hub#
📚 В этом руководстве объясняется, как загрузить YOLOv5 🚀 из PyTorch Hub по адресу https://pytorch.org/hub/ultralytics_yolov5.
Перед началом#
Установи 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. Подробную информацию обо всех доступных моделях см. в README.
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. 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
Все параметры инференса см. в методе AutoShape() forward для YOLOv5.
Параметры инференса#
Модели 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. В этом случае тебе нужно предоставить собственный скрипт обучения. Кроме того, см. наше руководство по обучению модели на пользовательских данных YOLOv5 — Train Custom Data Tutorial.
import torch
model = torch.hub.load("ultralytics/yolov5", "yolov5s", autoshape=False) # загрузить предварительно обученную модель
model = torch.hub.load("ultralytics/yolov5", "yolov5s", autoshape=False, pretrained=False) # загрузить модель с нуляРезультаты в Base64#
Для использования с API-сервисами. Подробности см. в примере Flask REST API.
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, включая модели, обученные на пользовательских данных. Подробнее об экспорте моделей см. в руководстве по экспорту в TFLite, ONNX, CoreML и TensorRT.
- 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.