跳至内容

PyTorch 集线器

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

开始之前

将requirements.txt安装在 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 😃

负载YOLOv5 与PyTorch Hub

简单示例

该示例从PyTorch Hub 中加载预训练的 YOLOv5s 模型，即 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  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, 显示 在支持的环境中进行筛选，并返回为 张量 或 大熊猫 数据帧。

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 齐达内的推理结果.jpg YOLO 公共汽车上的推理结果.jpg

有关所有推理选项，请参见YOLOv5 AutoShape() 转发 方法.

推理设置

YOLOv5 模型包含各种推理属性，如置信度阈值、IoU阈值等，这些属性可由参数设置：

model.conf = 0.25  # NMS confidence threshold
iou = 0.45  # NMS IoU threshold
agnostic = False  # NMS class-agnostic
multi_label = False  # NMS multiple labels per box
classes = None  # (optional list) filter by class, i.e. = [0, 15, 16] for COCO persons, cats and dogs
max_det = 1000  # maximum number of detections per image
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

输入通道

要加载一个预训练的 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 通过丢弃现有缓存并强制从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 torch
import threading


def run(model, im):
    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 服务。详见 https://github.com/ultralytics/yolov5/pull/2291 和Flask REST API示例。

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 DataFrames 返回：

results = model(im)  # inference
results.pandas().xyxy[0]  # Pandas DataFrame
熊猫输出（点击展开） 
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 结果

一旦转换为 .pandas() 数据帧使用 .to_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 级 挥发性有机化合物-训练有素的 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 导出教程。

💡 专业提示： TensorRT可能比PyTorch 快 2-5 倍 GPU 基准测试 💡 ProTip： ONNX和 OpenVINO可能比PyTorch 快 2-3 倍。 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.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 小时和每次新提交时都会进行一次测试。



创建于 2023-11-12，更新于 2023-12-03
作者:glenn-jocher(3)

评论