从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
有关所有推理选项,请参阅 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 小时进行一次,并在每次提交新内容时进行。