PyTorch Hub

📚 本指南解释了如何从PyTorch Hub加载YOLOv5 🚀，地址为https://pytorch.org/hub/ultralytics_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 = torch.hub.load("ultralytics/yolov5", "yolov5s")

# 图像
im = "https://ultralytics.com/images/zidane.jpg"

# 推理
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，在支持的环境中显示在屏幕上，并作为张量或pandas数据框返回。

import cv2
import torch
from PIL import Image

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

# 图像
for f in "zidane.jpg", "bus.jpg":
    torch.hub.download_url_to_file("https://ultralytics.com/images/" + f, f)  # 下载2张图像
im1 = Image.open("zidane.jpg")  # PIL图像
im2 = cv2.imread("bus.jpg")[..., ::-1]  # OpenCV图像（BGR到RGB）

# 推理
results = model([im1, im2], size=640)  # 图像批量

# 结果
results.print()
results.save()  # 或.show()

results.xyxy[0]  # im1预测（张量）
results.pandas().xyxy[0]  # im1预测（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置信度阈值
iou = 0.45  # NMS IoU阈值
agnostic = False  # NMS类不可知
multi_label = False  # NMS每个框多个标签
classes = None  # （可选列表）按类过滤，例如 = [0, 15, 16] 用于COCO中的人、猫和狗
max_det = 1000  # 每张图像的最大检测数
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)  # 静默加载

输入通道

要加载一个具有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)  # 强制重新加载

屏幕截图推理

要在桌面上运行推理：

import torch
from PIL import ImageGrab

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

# 图像
im = ImageGrab.grab()  # 截取屏幕截图

# 推理
results = model(im)

多GPU推理

YOLOv5模型可以加载到多个GPU上并行进行线程推理：

import threading

import torch


def run(model, im):
    """使用给定模型对图像进行推理并保存输出；模型必须支持`.save()`方法。"""
    results = model(im)
    results.save()


# 模型
model0 = torch.hub.load("ultralytics/yolov5", "yolov5s", device=0)
model1 = torch.hub.load("ultralytics/yolov5", "yolov5s", device=1)

# 推理
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)  # 加载预训练模型
model = torch.hub.load("ultralytics/yolov5", "yolov5s", autoshape=False, pretrained=False)  # 加载随机初始化模型

Base64结果

用于API服务。详情请参阅https://github.com/ultralytics/yolov5/pull/2291和Flask REST API示例。

results = model(im)  # 推理

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)  # 推理
crops = results.crop(save=True)  # 裁剪检测字典

Pandas结果

结果可以作为Pandas DataFrames返回：

results = model(im)  # 推理
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)  # 推理
results.pandas().xyxy[0].sort_values("xmin")  # 从左到右排序

框裁剪结果

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

results = model(im)  # 推理
crops = results.crop(save=True)  # 裁剪检测字典

JSON结果

结果可以在转换为.pandas()数据框后使用.to_json()方法以JSON格式返回。JSON格式可以使用orient参数进行修改。详情请参阅pandas .to_json()文档。

results = model(ims)  # 推理
results.pandas().xyxy[0].to_json(orient="records")  # JSON 图片1的预测结果

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")  # 本地模型
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 在GPU 基准测试上可能比 PyTorch 快 2-5 倍 💡 专业提示：ONNX 和 OpenVINO 在CPU 基准测试上可能比 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)
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 小时进行一次，并在每次新提交时进行。

📅 Created 11 months ago ✏️ Updated 20 days ago