在 GitHub 上编辑页面 (EN)

从 PyTorch Hub 加载 YOLOv5

📚 本指南介绍了如何从 PyTorch Hub 的 https://pytorch.org/hub/ultralytics_yolov5 加载 YOLOv5 🚀。

在你开始之前

Install requirements.txt in a Python>=3.8.0 environment, including PyTorch>=1.8. Models and datasets download automatically from the latest YOLOv5 release.

pip install -r https://raw.githubusercontent.com/ultralytics/yolov5/master/requirements.txt
提示

无需克隆 ultralytics/yolov5 —— PyTorch Hub 会自动获取代码。

使用 PyTorch Hub 加载 YOLOv5

简单示例

This example loads a pretrained YOLOv5s model from PyTorch Hub as model and passes an image for inference. 'yolov5s' is the lightest and fastest YOLOv5 model. For details on all available models please see the 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  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

详细示例

此示例展示了使用 PIL 和 OpenCV 图像源进行批量推理。results 可以打印到控制台，保存到 runs/hub，在支持的环境中显示在屏幕上，并以 tensors 或 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("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
YOLO inference results on zidane.jpg YOLO inference results on bus.jpg

有关所有推理选项，请参阅 YOLOv5 AutoShape() 的前向 方法。

推理设置

YOLOv5 模型包含各种推理属性，例如 置信度阈值、IoU 阈值等，可以通过以下方式设置：

model.conf = 0.25  # NMS confidence threshold
model.iou = 0.45  # NMS IoU threshold
model.agnostic = False  # NMS class-agnostic
model.multi_label = False  # NMS multiple labels per box
model.classes = None  # (optional list) filter by class, i.e. = [0, 15, 16] for COCO persons, cats and dogs
model.max_det = 1000  # maximum number of detections per image
model.amp = False  # Automatic Mixed Precision (AMP) inference

results = model(im, size=320)  # custom inference size

设备

模型创建后可以迁移到任何设备：

model.cpu()  # CPU
model.cuda()  # GPU
model.to(device)  # i.e. device=torch.device(0)

模型也可以直接在任何 device 上创建：

model = torch.hub.load("ultralytics/yolov5", "yolov5s", device="cpu")  # load on CPU
提示

输入图像在推理前会自动迁移到正确的模型设备。

静默输出

可以使用 _verbose=False 静默加载模型：

model = torch.hub.load("ultralytics/yolov5", "yolov5s", _verbose=False)  # load silently

输入通道

要加载一个具有 4 个输入通道（而非默认 3 个）的预训练 YOLOv5s 模型：

model = torch.hub.load("ultralytics/yolov5", "yolov5s", channels=4)

在这种情况下，模型将由预训练权重组成，除了最开头的输入层（其形状不再与预训练的输入层相同）。输入层将保持由随机权重初始化。

类别数量

要加载一个具有 10 个输出类别（而非默认 80 个）的预训练 YOLOv5s 模型：

model = torch.hub.load("ultralytics/yolov5", "yolov5s", classes=10)

在这种情况下，模型将由预训练权重组成，除了输出层（其形状不再与预训练的输出层相同）。输出层将保持由随机权重初始化。

强制重新加载

如果你在上述步骤中遇到问题，设置 force_reload=True 可能有帮助，它会丢弃现有的缓存并强制从 PyTorch Hub 重新下载最新版本的 YOLOv5。缓存副本位于 ~/.cache/torch/hub；删除该文件夹可达到相同的效果。

model = torch.hub.load("ultralytics/yolov5", "yolov5s", force_reload=True)  # force reload

屏幕截图推理

要在桌面屏幕上运行推理：

import torch
from PIL import ImageGrab

# Model
model = torch.hub.load("ultralytics/yolov5", "yolov5s")

# Image
im = ImageGrab.grab()  # take a screenshot

# 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 训练自定义数据教程 进行模型训练。

import torch

model = torch.hub.load("ultralytics/yolov5", "yolov5s", autoshape=False)  # load pretrained
model = torch.hub.load("ultralytics/yolov5", "yolov5s", autoshape=False, pretrained=False)  # load scratch

Base64 结果

用于 API 服务。详情请参阅 Flask REST API 示例。

import base64
from io import BytesIO

from PIL import Image

results = model(im)  # inference

results.ims  # array of original images (as np array) passed to model for inference
results.render()  # updates results.ims with boxes and labels
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 encoded image with results

裁剪结果

结果可以作为检测裁剪图返回并保存：

results = model(im)  # inference
crops = results.crop(save=True)  # cropped detections dictionary

Pandas 结果

结果可以作为 Pandas DataFrames 返回：

results = model(im)  # inference
results.pandas().xyxy[0]  # Pandas DataFrame
Pandas Output (click to expand) 
print(results.pandas().xyxy[0])
#      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

排序结果

结果可以按列排序，例如将车牌数字检测结果从左到右（x 轴）排序：

results = model(im)  # inference
results.pandas().xyxy[0].sort_values("xmin")  # sorted left-right

JSON 结果

一旦使用 .pandas() 方法转换为数据帧，结果就可以以 JSON 格式返回，使用 .to_json() 方法即可。JSON 格式可以使用 orient 参数进行修改。有关详情，请参阅 pandas .to_json() 文档。

results = model(ims)  # inference
results.pandas().xyxy[0].to_json(orient="records")  # JSON img1 predictions
JSON Output (click to expand) 
[
    {
        "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"
    }
]

自定义模型

This example loads a custom 20-class VOC-trained YOLOv5s model 'best.pt' with PyTorch Hub.

import torch

model = torch.hub.load("ultralytics/yolov5", "custom", path="path/to/best.pt")  # local model
model = torch.hub.load("path/to/yolov5", "custom", path="path/to/best.pt", source="local")  # local repo

TensorRT、ONNX 和 OpenVINO 模型

PyTorch Hub 支持大多数 YOLOv5 导出格式的推理，包括自定义训练的模型。有关导出模型的详细信息，请参阅 TFLite、ONNX、CoreML、TensorRT 导出教程。

提示
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.mlmodel")  # CoreML (macOS-only)
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 等关键依赖项，助你快速启动项目。

项目状态

YOLOv5 CI

此徽章表示所有 YOLOv5 GitHub Actions 持续集成 (CI) 测试均成功通过。这些 CI 测试严格检查 YOLOv5 在各个关键方面的功能和性能：训练、验证、推理、导出和基准测试。它们确保了在 macOS、Windows 和 Ubuntu 上的一致且可靠的运行，测试每 24 小时以及在每次新提交时进行。

Contributors
在 GitHub 上编辑页面 (EN)
GLglenn-jocher13PDpderrenger1RIRizwanMunawar1BUBurhan-Q1

评论