从PyTorch Hub加载YOLOv5

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

开始之前

在 Python>=3.8.0 环境中安装 requirements.txt，包括 PyTorch>=1.8。模型和数据集会自动从最新的 YOLOv5 版本下载。

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

💡 专家提示：不需要克隆 https://github.com/ultralytics/yolov5 😃

使用 PyTorch Hub 加载 YOLOv5

简单示例

此示例从 PyTorch Hub 加载一个预训练的 YOLOv5s 模型作为 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  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 在 zidane.jpg 上的推理结果 YOLO 在 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 版本可能会有所帮助。

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 输出 (点击展开)

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

框裁剪结果

结果可以返回并另存为检测裁剪：

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

JSON 结果

结果可以在转换为 JSON 格式后返回 .pandas() 使用以下项的数据帧 .to_json() method。可以使用以下方法修改 JSON 格式 orient 参数。见 pandas .to_json() 文档详情请见。

results = model(ims)  # inference
results.pandas().xyxy[0].to_json(orient="records")  # JSON img1 predictions

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"
    }
]

自定义模型

此示例加载一个自定义的 20 类 VOC预训练的 YOLOv5s 模型 'best.pt' 通过 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 导出教程。

💡 专家提示：在 GPU 基准测试 中，TensorRT 可能比 PyTorch 快 2-5 倍 💡 专家提示：在 CPU 基准测试 中，ONNX 和 OpenVINO 可能比 PyTorch 快 2-3 倍

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，以快速启动您的项目。

免费 GPU Notebooks:
Google Cloud：GCP 快速入门指南
Amazon：AWS 快速入门指南
Azure：AzureML 快速入门指南
Docker: Docker 快速入门指南

项目状态

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

📅 创建于 1 年前 ✏️ 更新于 5 个月前