yolort 直觉

[1]:

import cv2
import torch

from yolort.utils import Visualizer, get_image_from_url, read_image_to_tensor
from yolort.v5.utils.downloads import safe_download

[2]:

import os

os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"]="0"

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

读取图片

请求并预处理要检测的图像。

[3]:

img_source = "https://huggingface.co/spaces/zhiqwang/assets/resolve/main/bus.jpg"
# img_source = "https://huggingface.co/spaces/zhiqwang/assets/resolve/main/zidane.jpg"
img_raw = get_image_from_url(img_source)
img = read_image_to_tensor(img_raw)
img = img.to(device)

images = [img]

模型定义和初始化

[4]:

score_thresh = 0.55
stride = 64
img_h, img_w = 640, 640

[5]:

from yolort.models import yolov5n6

[6]:

model = yolov5n6(pretrained=True, score_thresh=score_thresh, size=(img_h, img_w))

Downloading: "https://github.com/zhiqwang/yolov5-rt-stack/releases/download/v0.5.2-alpha/yolov5_darknet_pan_n6_r60_coco-4e823e0f.pt" to /home/pc/.cache/torch/hub/checkpoints/yolov5_darknet_pan_n6_r60_coco-4e823e0f.pt

[7]:

model = model.eval()
model = model.to(device)

[8]:

# Perform inference on an image tensor
model_out = model(images)

/media/pc/data/4tb/lxw/anaconda3/envs/torch/lib/python3.10/site-packages/torch/functional.py:568: UserWarning: torch.meshgrid: in an upcoming release, it will be required to pass the indexing argument. (Triggered internally at  /opt/conda/conda-bld/pytorch_1646755897462/work/aten/src/ATen/native/TensorShape.cpp:2228.)
  return _VF.meshgrid(tensors, **kwargs)  # type: ignore[attr-defined]

验证 PyTorch 后端推理结果

[9]:

model_out[0]['boxes'].cpu().detach()

[9]:

tensor([[ 32.27846, 225.15266, 811.47729, 740.91071],
        [ 50.42179, 387.48911, 241.54399, 897.61035],
        [219.03331, 386.14352, 345.77689, 869.02582],
        [678.05023, 374.65326, 809.80341, 874.80621]])

（yolov5n6，形状为 1080x810）的目标输出：

tensor([[ 32.27846, 225.15259, 811.47729, 740.91077],
        [ 50.42178, 387.48911, 241.54393, 897.61035],
        [219.03334, 386.14346, 345.77686, 869.02582],
        [678.05023, 374.65341, 809.80341, 874.80621]])

[10]:

model_out[0]['scores'].cpu().detach()

[10]:

tensor([0.88238, 0.84486, 0.72629, 0.70077])

（yolov5n6，形状为 1080x810）的目标得分：

tensor([0.88238, 0.84486, 0.72629, 0.70077])

[11]:

model_out[0]['labels'].cpu().detach()

[11]:

tensor([5, 0, 0, 0])

目标输出的标签：

tensor([5, 0, 0, 0])

检测输出可视化

首先，获取 COCO 数据集的标签。

[12]:

# label_path = "https://raw.githubusercontent.com/zhiqwang/yolov5-rt-stack/main/notebooks/assets/coco.names"
label_source = "https://huggingface.co/spaces/zhiqwang/assets/resolve/main/coco.names"
label_path = label_source.split("/")[-1]
safe_download(label_path, label_source)

Downloading https://huggingface.co/spaces/zhiqwang/assets/resolve/main/coco.names to coco.names...



[13]:

v = Visualizer(img_raw, metalabels=label_path)
v.draw_instance_predictions(model_out[0])
v.imshow(scale=0.5)
../_images/notebooks_inference-pytorch-export-libtorch_20_0.png

脚本化 YOLOv5

[14]:

# TorchScript export
print(f'Starting TorchScript export with torch {torch.__version__}...')
export_script_name = 'yolov5n6.torchscript.pt'  # filename

Starting TorchScript export with torch 1.11.0...

[15]:

model_script = torch.jit.script(model)
model_script.eval()
model_script = model_script.to(device)

[16]:

# 保存脚本模型文件以供后续使用(可选)
model_script.save(export_script_name)

对 TorchScript 后端的推理

[17]:

out = model(images)
out_script = model_script(images)

/media/pc/data/4tb/lxw/anaconda3/envs/torch/lib/python3.10/site-packages/yolort/models/yolo.py:179: UserWarning: YOLO always returns a (Losses, Detections) tuple in scripting.
  warnings.warn("YOLO always returns a (Losses, Detections) tuple in scripting.")
/media/pc/data/4tb/lxw/anaconda3/envs/torch/lib/python3.10/site-packages/yolort/models/yolov5.py:180: UserWarning: YOLOv5 always returns a (Losses, Detections) tuple in scripting.
  warnings.warn("YOLOv5 always returns a (Losses, Detections) tuple in scripting.")

在 LibTorch 后端验证推理输出

[18]:

for k, v in out[0].items():
    torch.testing.assert_allclose(out_script[1][0][k], v, rtol=1e-07, atol=1e-09)

print("Exported model has been tested with libtorch, and the result looks good!")

Exported model has been tested with libtorch, and the result looks good!

比较 pytorch 和 libtorch 的推理时间

PyTorch 后端消耗的时间

[19]:

%%timeit
with torch.no_grad():
    out = model(images)

17.1 ms ± 54.2 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

LibTorch 后端消耗的时间

[20]:

# Dummy run once for torchscript
with torch.no_grad():
    out_script = model_script(images)

[21]:

%%timeit
with torch.no_grad():
    out_script = model_script(images)

11.3 ms ± 430 µs per loop (mean ± std. dev. of 7 runs, 1 loop each)

View this document as a notebook: https://github.com/zhiqwang/yolov5-rt-stack/blob/main/notebooks/inference-pytorch-export-libtorch.ipynb