autobackend.py
ultralytics\nn\autobackend.py
目录
autobackend.py
1.所需的库和模块
2.def check_class_names(names):
3.def default_class_names(data=None):
4.class AutoBackend(nn.Module):
1.所需的库和模块
# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
import ast
import json
import platform
import zipfile
from collections import OrderedDict, namedtuple
from pathlib import Path
import cv2
import numpy as np
import torch
import torch.nn as nn
from PIL import Image
from ultralytics.utils import ARM64, IS_JETSON, IS_RASPBERRYPI, LINUX, LOGGER, PYTHON_VERSION, ROOT, yaml_load
from ultralytics.utils.checks import check_requirements, check_suffix, check_version, check_yaml, is_rockchip
from ultralytics.utils.downloads import attempt_download_asset, is_url2.def check_class_names(names):
# 这段代码定义了一个函数 check_class_names ，用于检查和处理类别名称（ names ）。它的主要功能包。将类别名称从列表转换为字典。将类别索引从字符串转换为整数，并确保类别名称为字符串格式。检查类别索引是否有效。如果类别名称是 ImageNet 的类别代码，则将其映射为人类可读的名称。
# 定义了一个函数 check_class_names ，接受一个参数。
# 1.names ：类别名称。
def check_class_names(names):
# 检查类名。
# 如果需要，将 imagenet 类代码映射到人类可读的名称。将列表转换为字典。
"""
Check class names.
Map imagenet class codes to human-readable names if required. Convert lists to dicts.
"""
# 检查 names 是否是一个列表。
if isinstance(names, list): # names is a list
# 如果 names 是一个列表，则将其转换为字典。 使用 enumerate(names) 生成索引和值的元组。 使用 dict(...) 将其转换为字典，其中索引作为键，类别名称作为值。
names = dict(enumerate(names)) # convert to dict
# 检查 names 是否是一个字典。
if isinstance(names, dict):
# Convert 1) string keys to int, i.e. '0' to 0, and non-string values to strings, i.e. True to 'True'
# 如果 names 是一个字典，则进行以下处理。
# 将键（类别索引）从字符串转换为整数。
# 将值（类别名称）转换为字符串。
# 使用字典推导式生成新的字典。
names = {int(k): str(v) for k, v in names.items()}
# 获取 类别总数 ，存储在变量 n 中。
n = len(names)
# 检查 类别索引是否有效 。如果最大类别索引大于或等于类别总数 n ，则说明类别索引无效。
if max(names.keys()) >= n:
# 抛出一个 KeyError 异常，提示用户类别索引无效。
raise KeyError(
f"{n}-class dataset requires class indices 0-{n - 1}, but you have invalid class indices " # {n} 类数据集需要类索引 0-{n - 1}，
f"{min(names.keys())}-{max(names.keys())} defined in your dataset YAML." # 但是您在数据集 YAML 中定义了无效的类索引 {min(names.keys())}-{max(names.keys())}。
)
# 检查类别名称是否为 ImageNet 的类别代码。 如果第一个类别名称是字符串，并且以 "n0" 开头（例如 "n01440764" ），则认为是 ImageNet 的类别代码。
if isinstance(names[0], str) and names[0].startswith("n0"): # imagenet class codes, i.e. 'n01440764'
# 加载 ImageNet 的类别映射文件。 使用 yaml_load 函数加载 ROOT / "cfg/datasets/ImageNet.yaml" 文件。 提取文件中的 map 字段，该字段包含类别代码到人类可读名称的映射。
# def yaml_load(file="data.yaml", append_filename=False): -> 用于加载和解析YAML文件。它还提供了一些额外的功能，例如清理文件内容中的特殊字符，并可以选择将文件名添加到解析后的数据中。返回解析后的YAML数据，可能包含文件名（如果 append_filename 为 True ）。 -> return data
names_map = yaml_load(ROOT / "cfg/datasets/ImageNet.yaml")["map"] # human-readable names
# 将类别代码映射为人类可读的名称。 使用字典推导式，将 names 中的每个类别代码 v 替换为对应的可读名称 names_map[v] 。
names = {k: names_map[v] for k, v in names.items()}
# 返回处理后的类别名称字典。
return names
# check_class_names 函数的主要功能和作用如下。格式转换：将类别名称从列表转换为字典，方便后续处理。数据类型校验：确保类别索引为整数，类别名称为字符串。索引检查：验证类别索引是否在有效范围内，避免索引错误。映射处理：如果类别名称是 ImageNet 的类别代码，则将其映射为人类可读的名称。灵活性和兼容性：支持多种输入格式（列表或字典），并提供详细的错误提示，方便调试和使用。这个函数在处理类别名称时非常实用，尤其是在处理数据集配置文件时，能够确保类别名称的正确性和一致性。3.def default_class_names(data=None):
# 这段代码定义了一个函数 default_class_names ，用于从输入的 YAML 文件中提取类别名称，或者在无法提取时返回默认的类别名称。
# 定义了一个函数 default_class_names ，接受一个可选参数。
# 1.data ：输入的 YAML 文件路径或其他数据源。
def default_class_names(data=None):
# 将默认类名应用于输入 YAML 文件或返回数字类名。
"""Applies default class names to an input YAML file or returns numerical class names."""
# 检查是否提供了输入数据 data 。
if data:
# 开始一个 try 块，尝试从输入数据中提取类别名称。
try:
# 调用 check_yaml(data) 函数，验证输入的 data 是否为有效的 YAML 文件路径或内容。
# 使用 yaml_load 函数加载 YAML 文件内容。
# 提取 names 字段，返回类别名称列表。
# def check_yaml(file, suffix=(".yaml", ".yml"), hard=True): -> 用于检查YAML文件是否存在，如果不存在则尝试下载文件，并返回文件的路径。它还检查文件的后缀是否符合指定的允许后缀列表。调用 check_file 函数，检查文件是否存在，如果不存在则尝试下载文件，并返回文件的路径。 -> return check_file(file, suffix, hard=hard)
return yaml_load(check_yaml(data))["names"]
# 如果在尝试提取类别名称时发生任何异常（如文件不存在、格式错误等），则进入 except 块。
except Exception:
# 在 except 块中，不执行任何操作，直接跳过异常处理。
pass
# 如果无法从输入数据中提取类别名称（或未提供输入），则返回默认的类别名称字典。 使用字典推导式生成一个字典，键为类别索引（从 0 到 998），值为格式化的字符串 "class{i}" 。 这样可以确保即使没有提供有效的类别名称，函数也能返回一个默认的类别名称映射。
return {i: f"class{i}" for i in range(999)} # return default if above errors
# default_class_names 函数的主要功能和作用如下。从 YAML 文件中提取类别名称：如果提供了有效的 YAML 文件路径或内容，函数会尝试从中提取 names 字段。提取的类别名称以列表形式返回。提供默认的类别名称：如果输入数据无效或未提供输入，函数会返回一个默认的类别名称字典。默认的类别名称格式为 "class{i}" ，其中 i 是类别索引。容错处理：函数通过 try-except 块捕获可能的异常，确保即使输入数据有问题，也不会导致程序崩溃。如果提取失败，函数会返回默认的类别名称字典，保证函数的输出始终是有效的。这个函数在处理数据集配置文件时非常实用，尤其是在类别名称可能缺失或格式不正确的情况下，能够提供一个可靠的默认值。4.class AutoBackend(nn.Module):
# 这段代码定义了一个名为 AutoBackend 的类，它是 Ultralytics YOLO 模型的动态后端选择器，用于在不同平台上运行推理。 AutoBackend 支持多种模型格式，并能够根据输入模型的格式动态切换后端，从而简化了模型部署过程。
# AutoBackend 继承自 torch.nn.Module ，使其可以像 PyTorch 模型一样使用。
class AutoBackend(nn.Module):
# 处理使用 Ultralytics YOLO 模型运行推理的动态后端选择。
# AutoBackend 类旨在为各种推理引擎提供抽象层。它支持多种格式，每种格式都有特定的命名约定，如下所述：
# 支持的格式和命名约定：
# | Format | File Suffix |
# | --------------------- | ----------------- |
# | PyTorch | *.pt |
"""
Handles dynamic backend selection for running inference using Ultralytics YOLO models.
The AutoBackend class is designed to provide an abstraction layer for various inference engines. It supports a wide
range of formats, each with specific naming conventions as outlined below:
Supported Formats and Naming Conventions:
| Format | File Suffix |
| --------------------- | ----------------- |
| PyTorch | *.pt |
| TorchScript | *.torchscript |
| ONNX Runtime | *.onnx |
| ONNX OpenCV DNN | *.onnx (dnn=True) |
| OpenVINO | *openvino_model/ |
| CoreML | *.mlpackage |
| TensorRT | *.engine |
| TensorFlow SavedModel | *_saved_model/ |
| TensorFlow GraphDef | *.pb |
| TensorFlow Lite | *.tflite |
| TensorFlow Edge TPU | *_edgetpu.tflite |
| PaddlePaddle | *_paddle_model/ |
| MNN | *.mnn |
| NCNN | *_ncnn_model/ |
| IMX | *_imx_model/ |
| RKNN | *_rknn_model/ |
This class offers dynamic backend switching capabilities based on the input model format, making it easier to deploy
models across various platforms.
"""
# 这段代码是 AutoBackend 类的初始化方法 __init__ ，它负责根据输入的模型权重和配置参数，加载不同格式的模型，并进行相应的初始化操作。
# 装饰器，表示在初始化过程中不会计算梯度，适用于推理阶段。
@torch.no_grad()
# 参数 ：
# 1.weights ：模型权重文件路径或 torch.nn.Module 实例，默认为 "yolo11n.pt" 。
# 2.device ：运行模型的设备，默认为 CPU。
# 3.dnn ：是否使用 OpenCV DNN 模块进行 ONNX 推理，默认为 False 。
# 4.data ：包含类别名称的额外 data.yaml 文件路径，可选。
# 5.fp16 ：是否启用半精度推理，默认为 False 。
# 6.batch ：推理时假设的批量大小，默认为 1。
# 7.fuse ：是否融合 Conv2D + BatchNorm 层以优化，默认为 True 。
# 8.verbose ：是否启用详细日志，默认为 True 。
def __init__(
self,
weights="yolo11n.pt",
device=torch.device("cpu"),
dnn=False,
data=None,
fp16=False,
batch=1,
fuse=True,
verbose=True,
):
# 初始化 AutoBackend 进行推理。
"""
Initialize the AutoBackend for inference.
Args:
weights (str | torch.nn.Module): Path to the model weights file or a module instance. Defaults to 'yolo11n.pt'.
device (torch.device): Device to run the model on. Defaults to CPU.
dnn (bool): Use OpenCV DNN module for ONNX inference. Defaults to False.
data (str | Path | optional): Path to the additional data.yaml file containing class names. Optional.
fp16 (bool): Enable half-precision inference. Supported only on specific backends. Defaults to False.
batch (int): Batch-size to assume for inference.
fuse (bool): Fuse Conv2D + BatchNorm layers for optimization. Defaults to True.
verbose (bool): Enable verbose logging. Defaults to True.
"""
# 这段代码是 AutoBackend 类初始化方法 __init__ 的一部分，主要负责初始化一些基础变量和模型的格式检测。
# 调用父类 torch.nn.Module 的初始化方法，确保 AutoBackend 类继承了 PyTorch 模块的基本功能。
super().__init__()
# 如果 weights 是一个列表，则取列表的第一个元素作为权重路径；否则直接将 weights 转换为字符串。 这样可以支持用户传入单个权重文件路径或权重文件列表。
w = str(weights[0] if isinstance(weights, list) else weights)
# 检查 weights 是否为 torch.nn.Module 的实例。 如果是，说明用户传入的是一个已经加载的 PyTorch 模型，而不是权重文件路径。
nn_module = isinstance(weights, torch.nn.Module)
# 检测模型格式。
(
# _model_type 方法返回一个元组，包含布尔值。
# 表示模型是否为 PyTorch ( pt )、
pt,
# TorchScript ( jit )、
jit,
# ONNX ( onnx )、
onnx,
# OpenVINO ( xml )、
xml,
# TensorRT ( engine )、
engine,
# CoreML ( coreml )、
coreml,
# TensorFlow SavedModel ( saved_model )、
saved_model,
# TensorFlow GraphDef ( pb )、
pb,
# TensorFlow Lite ( tflite )、
tflite,
# Edge TPU ( edgetpu )、
edgetpu,
# TensorFlow.js ( tfjs )、
tfjs,
# PaddlePaddle ( paddle )、
paddle,
# MNN ( mnn )、
mnn,
# NCNN ( ncnn )、
ncnn,
# IMX ( imx )、
imx,
# RKNN ( rknn ) 、
rknn,
# 或 Triton ( triton ) 格式。
triton,
# 调用 _model_type 方法，根据权重文件路径 w 检测模型的格式。
) = self._model_type(w)
# 设置 FP16 和 NHWC 格式。
# 如果模型是 PyTorch ( pt )、TorchScript ( jit )、ONNX ( onnx )、OpenVINO ( xml )、TensorRT ( engine )、PyTorch 模块 ( nn_module ) 或 Triton ( triton ) 格式，则支持半精度推理（FP16）。
# 这里使用了位运算符 &= ，只有在上述条件满足时， fp16 才保持为 True 。
fp16 &= pt or jit or onnx or xml or engine or nn_module or triton # FP16
# 如果模型是 CoreML ( coreml )、TensorFlow SavedModel ( saved_model )、TensorFlow GraphDef ( pb )、TensorFlow Lite ( tflite )、Edge TPU ( edgetpu ) 或 RKNN ( rknn ) 格式，则数据格式为 NHWC（ height, width, channels ），而不是 PyTorch 默认的 NCHW（ channels, height, width ）。
nhwc = coreml or saved_model or pb or tflite or edgetpu or rknn # BHWC formats (vs torch BCWH)
# 初始化其他变量。
# stride 默认步幅为 32。
stride = 32 # default stride
# end2end 默认为 False ，表示模型不是端到端的。
end2end = False # default end2end
# model 、 metadata 和 task 初始化为 None ，稍后根据模型格式加载相应的值。
model, metadata, task = None, None, None
# 这段代码的主要作用是。初始化基础变量：设置一些默认值，如步幅、端到端标志等。检测模型格式：根据权重文件路径或 PyTorch 模块，确定模型的格式。设置推理配置：根据模型格式，决定是否支持半精度推理（FP16）和数据格式（NHWC 或 NCHW）。这些初始化操作为后续的模型加载和推理提供了必要的基础信息。
# 这段代码的作用是根据模型的格式和设备的可用性来设置推理时使用的设备（CPU 或 GPU）。
# Set device
# 如果系统支持 CUDA 且用户指定的设备不是 CPU，则将 cuda 设置为 True ，表示可以使用 GPU 进行推理；否则， cuda 为 False 。
# torch.cuda.is_available() ：检查当前系统是否支持 CUDA（即是否有可用的 NVIDIA GPU）。
# device.type != "cpu" ：检查用户指定的设备是否不是 CPU。
cuda = torch.cuda.is_available() and device.type != "cpu" # use CUDA
# 如果模型不是上述支持 GPU 的格式，则将设备设置为 CPU，并将 cuda 设置为 False 。
# cuda ：表示前面检测到系统支持 CUDA。
# not any([nn_module, pt, jit, engine, onnx, paddle]) ：检查模型是否不是以下格式 ：
# nn_module ：PyTorch 模块。
# pt ：PyTorch 模型。
# jit ：TorchScript 模型。
# engine ：TensorRT 模型。
# onnx ：ONNX 模型。
# paddle ：PaddlePaddle 模型。
if cuda and not any([nn_module, pt, jit, engine, onnx, paddle]): # GPU dataloader formats
device = torch.device("cpu")
cuda = False
# 这段代码的主要作用是。检查 CUDA 是否可用：如果系统支持 CUDA 且用户指定的设备不是 CPU，则尝试使用 GPU。检查模型格式是否支持 GPU：如果模型格式不支持 GPU 推理，则将设备切换到 CPU，并禁用 CUDA。通过这段代码， AutoBackend 类能够根据模型格式和设备的可用性，动态选择最适合的推理设备。
# Download if not local
# 如果模型不是 PyTorch ( pt )、Triton ( triton ) 或 torch.nn.Module ( nn_module ) 格式。
if not (pt or triton or nn_module):
# 则调用 attempt_download_asset 函数。 attempt_download_asset 函数会尝试下载模型文件，确保模型文件在本地可用。
# def attempt_download_asset(file, repo="ultralytics/assets", release="v8.3.0", **kwargs): -> 用于尝试下载指定的资产文件（如模型文件）。它会检查文件是否已存在，如果不存在，则尝试从指定的URL或GitHub仓库下载。返回文件的路径。 -> return str(file)
w = attempt_download_asset(w)
# 这段代码处理了当输入的 weights 是一个已经加载的 PyTorch 模型（ torch.nn.Module ）时的情况。
# In-memory PyTorch model
# nn_module 是之前通过 isinstance(weights, torch.nn.Module) 检测得到的布尔值，表示 weights 是否是一个 PyTorch 模型实例。 如果是 PyTorch 模型实例，则进入这个分支。
if nn_module:
# 将输入的 PyTorch 模型 weights 移动到指定的设备（ device ）上。 这一步确保模型的参数和缓冲区在正确的设备上，例如 CPU 或 GPU。
model = weights.to(device)
# 如果 fuse 参数为 True 。
if fuse:
# 则调用模型的 fuse 方法进行融合优化。
# 融合优化通常包括将多个层（如 Conv2D + BatchNorm ）合并为一个层，以提高推理速度。 verbose 参数控制是否输出融合过程的详细信息。
model = model.fuse(verbose=verbose)
# 获取模型的属性。
# 检查模型是否有 kpt_shape 属性。
if hasattr(model, "kpt_shape"):
# 如果有，则将其值赋给变量 kpt_shape 。 kpt_shape 通常用于姿势检测模型，表示关键点的形状。
kpt_shape = model.kpt_shape # pose-only
# 获取模型的步幅（ stride ），并确保其值至少为 32。 确保模型的步幅符合要求，避免过小的步幅导致推理问题。
stride = max(int(model.stride.max()), 32) # model stride
# 获取模型的类别名称（ names ）。 如果模型是一个包装模块（ model.module ），则从包装模块中获取类别名称；否则直接从模型中获取。 确保能够正确获取模型的类别名称，用于后续的推理和结果解析。
names = model.module.names if hasattr(model, "module") else model.names # get class names
# 设置模型的数据类型。
# 如果 fp16 参数为 True ，则将模型的数据类型转换为半精度浮点数（ float16 ）；否则保持为单精度浮点数（ float32 ）。 半精度浮点数可以减少模型的内存占用，并可能提高推理速度，但可能会牺牲一些精度。
model.half() if fp16 else model.float()
# 显式分配模型。将处理后的模型显式分配给类的 model 属性。 这样可以在类的其他方法中方便地调用模型，并进行设备转换（如 to() ）、数据类型转换（如 half() ）等操作。
self.model = model # explicitly assign for to(), cpu(), cuda(), half()
# 将变量 pt 设置为 True ，表示当前处理的是 PyTorch 模型。 这个变量在后续的代码中用于判断模型的类型，并进行相应的处理。
pt = True
# 这段代码的主要作用是。处理 PyTorch 模型实例：将输入的 PyTorch 模型移动到指定设备，并进行融合优化。获取模型的属性：提取模型的步幅、类别名称等重要属性。设置模型的数据类型：根据 fp16 参数，将模型转换为半精度或单精度浮点数。显式分配模型：将处理后的模型分配给类的属性，以便后续使用。通过这段代码， AutoBackend 类能够正确处理输入的 PyTorch 模型，并为其后续的推理做好准备。
# 这段代码处理了当输入的 weights 是一个 PyTorch 模型文件路径时的情况。
# PyTorch
# pt 是之前通过 _model_type 方法检测得到的布尔值，表示 weights 是否为 PyTorch 模型文件路径。 如果是 PyTorch 模型文件路径，则进入这个分支。
elif pt:
# 导入 attempt_load_weights 函数，该函数用于加载 PyTorch 模型权重。
# def attempt_load_weights(weights, device=None, inplace=True, fuse=False): -> 用于加载模型权重，支持加载单个模型或多个模型组成的集成模型（ensemble）。如果 ensemble 中只有一个模型。返回最终的模型集合 ensemble 。 -> return ensemble[-1] / return ensemble
from ultralytics.nn.tasks import attempt_load_weights
# 调用 attempt_load_weights 函数加载 PyTorch 模型。
model = attempt_load_weights(
# 如果 weights 是一个列表，则取列表的第一个元素作为权重路径；否则直接使用 w （之前已经将 weights 转换为字符串路径）。
# device=device ：指定将模型加载到的设备。
# inplace=True ：表示在加载模型时进行原地操作，以节省内存。
# fuse=fuse ：如果 fuse 参数为 True ，则在加载模型时进行融合优化。
weights if isinstance(weights, list) else w, device=device, inplace=True, fuse=fuse
)
# 获取模型的属性。
# 检查模型是否有 kpt_shape 属性。如果有，则将其值赋给变量 kpt_shape 。 kpt_shape 通常用于姿势检测模型，表示 关键点的形状 。
if hasattr(model, "kpt_shape"):
kpt_shape = model.kpt_shape # pose-only
# 获取模型的 步幅 （ stride ），并确保其值至少为 32。 为了确保模型的步幅符合要求，避免过小的步幅导致推理问题。
stride = max(int(model.stride.max()), 32) # model stride
# 获取模型的类别名称（ names ）。 如果模型是一个包装模块（ model.module ），则从包装模块中获取类别名称；否则直接从模型中获取。 确保能够正确获取模型的类别名称，用于后续的推理和结果解析。
names = model.module.names if hasattr(model, "module") else model.names # get class names
# 如果 fp16 参数为 True ，则将模型的数据类型转换为半精度浮点数（ float16 ）；否则保持为单精度浮点数（ float32 ）。 半精度浮点数可以减少模型的内存占用，并可能提高推理速度，但可能会牺牲一些精度。
model.half() if fp16 else model.float()
# 显式分配模型。
# 将处理后的模型显式分配给类的 model 属性。
# 这样可以在类的其他方法中方便地调用模型，并进行设备转换（如 to() ）、数据类型转换（如 half() ）等操作。
self.model = model # explicitly assign for to(), cpu(), cuda(), half()
# 这段代码的主要作用是。加载 PyTorch 模型：使用 attempt_load_weights 函数加载 PyTorch 模型文件，并根据 fuse 参数进行融合优化。获取模型的属性：提取模型的步幅、类别名称等重要属性。设置模型的数据类型：根据 fp16 参数，将模型转换为半精度或单精度浮点数。显式分配模型：将处理后的模型分配给类的属性，以便后续使用。通过这段代码， AutoBackend 类能够正确处理输入的 PyTorch 模型文件，并为其后续的推理做好准备。
# 可忽略-----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------↓
# TorchScript
elif jit:
LOGGER.info(f"Loading {w} for TorchScript inference...")
extra_files = {"config.txt": ""} # model metadata
model = torch.jit.load(w, _extra_files=extra_files, map_location=device)
model.half() if fp16 else model.float()
if extra_files["config.txt"]: # load metadata dict
metadata = json.loads(extra_files["config.txt"], object_hook=lambda x: dict(x.items()))
# ONNX OpenCV DNN
elif dnn:
LOGGER.info(f"Loading {w} for ONNX OpenCV DNN inference...")
check_requirements("opencv-python>=4.5.4")
net = cv2.dnn.readNetFromONNX(w)
# ONNX Runtime and IMX
elif onnx or imx:
LOGGER.info(f"Loading {w} for ONNX Runtime inference...")
check_requirements(("onnx", "onnxruntime-gpu" if cuda else "onnxruntime"))
if IS_RASPBERRYPI or IS_JETSON:
# Fix 'numpy.linalg._umath_linalg' has no attribute '_ilp64' for TF SavedModel on RPi and Jetson
check_requirements("numpy==1.23.5")
import onnxruntime
providers = ["CPUExecutionProvider"]
if cuda and "CUDAExecutionProvider" in onnxruntime.get_available_providers():
providers.insert(0, "CUDAExecutionProvider")
elif cuda: # Only log warning if CUDA was requested but unavailable
LOGGER.warning("WARNING ⚠️ Failed to start ONNX Runtime with CUDA. Using CPU...")
device = torch.device("cpu")
cuda = False
LOGGER.info(f"Using ONNX Runtime {providers[0]}")
if onnx:
session = onnxruntime.InferenceSession(w, providers=providers)
else:
check_requirements(
["model-compression-toolkit==2.1.1", "sony-custom-layers[torch]==0.2.0", "onnxruntime-extensions"]
)
w = next(Path(w).glob("*.onnx"))
LOGGER.info(f"Loading {w} for ONNX IMX inference...")
import mct_quantizers as mctq
from sony_custom_layers.pytorch.object_detection import nms_ort # noqa
session = onnxruntime.InferenceSession(
w, mctq.get_ort_session_options(), providers=["CPUExecutionProvider"]
)
task = "detect"
output_names = [x.name for x in session.get_outputs()]
metadata = session.get_modelmeta().custom_metadata_map
dynamic = isinstance(session.get_outputs()[0].shape[0], str)
fp16 = True if "float16" in session.get_inputs()[0].type else False
if not dynamic:
io = session.io_binding()
bindings = []
for output in session.get_outputs():
out_fp16 = "float16" in output.type
y_tensor = torch.empty(output.shape, dtype=torch.float16 if out_fp16 else torch.float32).to(device)
io.bind_output(
name=output.name,
device_type=device.type,
device_id=device.index if cuda else 0,
element_type=np.float16 if out_fp16 else np.float32,
shape=tuple(y_tensor.shape),
buffer_ptr=y_tensor.data_ptr(),
)
bindings.append(y_tensor)
# OpenVINO
elif xml:
LOGGER.info(f"Loading {w} for OpenVINO inference...")
check_requirements("openvino>=2024.0.0")
import openvino as ov
core = ov.Core()
w = Path(w)
if not w.is_file(): # if not *.xml
w = next(w.glob("*.xml")) # get *.xml file from *_openvino_model dir
ov_model = core.read_model(model=str(w), weights=w.with_suffix(".bin"))
if ov_model.get_parameters()[0].get_layout().empty:
ov_model.get_parameters()[0].set_layout(ov.Layout("NCHW"))
# OpenVINO inference modes are 'LATENCY', 'THROUGHPUT' (not recommended), or 'CUMULATIVE_THROUGHPUT'
inference_mode = "CUMULATIVE_THROUGHPUT" if batch > 1 else "LATENCY"
LOGGER.info(f"Using OpenVINO {inference_mode} mode for batch={batch} inference...")
ov_compiled_model = core.compile_model(
ov_model,
device_name="AUTO", # AUTO selects best available device, do not modify
config={"PERFORMANCE_HINT": inference_mode},
)
input_name = ov_compiled_model.input().get_any_name()
metadata = w.parent / "metadata.yaml"
# TensorRT
elif engine:
LOGGER.info(f"Loading {w} for TensorRT inference...")
if IS_JETSON and PYTHON_VERSION <= "3.8.0":
# fix error: `np.bool` was a deprecated alias for the builtin `bool` for JetPack 4 with Python <= 3.8.0
check_requirements("numpy==1.23.5")
try:
import tensorrt as trt # noqa https://developer.nvidia.com/nvidia-tensorrt-download
except ImportError:
if LINUX:
check_requirements("tensorrt>7.0.0,!=10.1.0")
import tensorrt as trt # noqa
check_version(trt.__version__, ">=7.0.0", hard=True)
check_version(trt.__version__, "!=10.1.0", msg="https://github.com/ultralytics/ultralytics/pull/14239")
if device.type == "cpu":
device = torch.device("cuda:0")
Binding = namedtuple("Binding", ("name", "dtype", "shape", "data", "ptr"))
logger = trt.Logger(trt.Logger.INFO)
# Read file
with open(w, "rb") as f, trt.Runtime(logger) as runtime:
try:
meta_len = int.from_bytes(f.read(4), byteorder="little") # read metadata length
metadata = json.loads(f.read(meta_len).decode("utf-8")) # read metadata
except UnicodeDecodeError:
f.seek(0) # engine file may lack embedded Ultralytics metadata
model = runtime.deserialize_cuda_engine(f.read()) # read engine
# Model context
try:
context = model.create_execution_context()
except Exception as e: # model is None
LOGGER.error(f"ERROR: TensorRT model exported with a different version than {trt.__version__}\n")
raise e
bindings = OrderedDict()
output_names = []
fp16 = False # default updated below
dynamic = False
is_trt10 = not hasattr(model, "num_bindings")
num = range(model.num_io_tensors) if is_trt10 else range(model.num_bindings)
for i in num:
if is_trt10:
name = model.get_tensor_name(i)
dtype = trt.nptype(model.get_tensor_dtype(name))
is_input = model.get_tensor_mode(name) == trt.TensorIOMode.INPUT
if is_input:
if -1 in tuple(model.get_tensor_shape(name)):
dynamic = True
context.set_input_shape(name, tuple(model.get_tensor_profile_shape(name, 0)[1]))
if dtype == np.float16:
fp16 = True
else:
output_names.append(name)
shape = tuple(context.get_tensor_shape(name))
else: # TensorRT < 10.0
name = model.get_binding_name(i)
dtype = trt.nptype(model.get_binding_dtype(i))
is_input = model.binding_is_input(i)
if model.binding_is_input(i):
if -1 in tuple(model.get_binding_shape(i)): # dynamic
dynamic = True
context.set_binding_shape(i, tuple(model.get_profile_shape(0, i)[1]))
if dtype == np.float16:
fp16 = True
else:
output_names.append(name)
shape = tuple(context.get_binding_shape(i))
im = torch.from_numpy(np.empty(shape, dtype=dtype)).to(device)
bindings[name] = Binding(name, dtype, shape, im, int(im.data_ptr()))
binding_addrs = OrderedDict((n, d.ptr) for n, d in bindings.items())
batch_size = bindings["images"].shape[0] # if dynamic, this is instead max batch size
# CoreML
elif coreml:
LOGGER.info(f"Loading {w} for CoreML inference...")
import coremltools as ct
model = ct.models.MLModel(w)
metadata = dict(model.user_defined_metadata)
# TF SavedModel
elif saved_model:
LOGGER.info(f"Loading {w} for TensorFlow SavedModel inference...")
import tensorflow as tf
keras = False # assume TF1 saved_model
model = tf.keras.models.load_model(w) if keras else tf.saved_model.load(w)
metadata = Path(w) / "metadata.yaml"
# TF GraphDef
elif pb: # https://www.tensorflow.org/guide/migrate#a_graphpb_or_graphpbtxt
LOGGER.info(f"Loading {w} for TensorFlow GraphDef inference...")
import tensorflow as tf
from ultralytics.engine.exporter import gd_outputs
def wrap_frozen_graph(gd, inputs, outputs):
"""Wrap frozen graphs for deployment."""
x = tf.compat.v1.wrap_function(lambda: tf.compat.v1.import_graph_def(gd, name=""), []) # wrapped
ge = x.graph.as_graph_element
return x.prune(tf.nest.map_structure(ge, inputs), tf.nest.map_structure(ge, outputs))
gd = tf.Graph().as_graph_def() # TF GraphDef
with open(w, "rb") as f:
gd.ParseFromString(f.read())
frozen_func = wrap_frozen_graph(gd, inputs="x:0", outputs=gd_outputs(gd))
try: # find metadata in SavedModel alongside GraphDef
metadata = next(Path(w).resolve().parent.rglob(f"{Path(w).stem}_saved_model*/metadata.yaml"))
except StopIteration:
pass
# TFLite or TFLite Edge TPU
elif tflite or edgetpu: # https://www.tensorflow.org/lite/guide/python#install_tensorflow_lite_for_python
try: # https://coral.ai/docs/edgetpu/tflite-python/#update-existing-tf-lite-code-for-the-edge-tpu
from tflite_runtime.interpreter import Interpreter, load_delegate
except ImportError:
import tensorflow as tf
Interpreter, load_delegate = tf.lite.Interpreter, tf.lite.experimental.load_delegate
if edgetpu: # TF Edge TPU https://coral.ai/software/#edgetpu-runtime
device = device[3:] if str(device).startswith("tpu") else ":0"
LOGGER.info(f"Loading {w} on device {device[1:]} for TensorFlow Lite Edge TPU inference...")
delegate = {"Linux": "libedgetpu.so.1", "Darwin": "libedgetpu.1.dylib", "Windows": "edgetpu.dll"}[
platform.system()
]
interpreter = Interpreter(
model_path=w,
experimental_delegates=[load_delegate(delegate, options={"device": device})],
)
device = "cpu" # Required, otherwise PyTorch will try to use the wrong device
else: # TFLite
LOGGER.info(f"Loading {w} for TensorFlow Lite inference...")
interpreter = Interpreter(model_path=w) # load TFLite model
interpreter.allocate_tensors() # allocate
input_details = interpreter.get_input_details() # inputs
output_details = interpreter.get_output_details() # outputs
# Load metadata
try:
with zipfile.ZipFile(w, "r") as model:
meta_file = model.namelist()[0]
metadata = ast.literal_eval(model.read(meta_file).decode("utf-8"))
except zipfile.BadZipFile:
pass
# TF.js
elif tfjs:
raise NotImplementedError("YOLOv8 TF.js inference is not currently supported.")
# PaddlePaddle
elif paddle:
LOGGER.info(f"Loading {w} for PaddlePaddle inference...")
check_requirements("paddlepaddle-gpu" if cuda else "paddlepaddle")
import paddle.inference as pdi # noqa
w = Path(w)
if not w.is_file(): # if not *.pdmodel
w = next(w.rglob("*.pdmodel")) # get *.pdmodel file from *_paddle_model dir
config = pdi.Config(str(w), str(w.with_suffix(".pdiparams")))
if cuda:
config.enable_use_gpu(memory_pool_init_size_mb=2048, device_id=0)
predictor = pdi.create_predictor(config)
input_handle = predictor.get_input_handle(predictor.get_input_names()[0])
output_names = predictor.get_output_names()
metadata = w.parents[1] / "metadata.yaml"
# MNN
elif mnn:
LOGGER.info(f"Loading {w} for MNN inference...")
check_requirements("MNN") # requires MNN
import os
import MNN
config = {"precision": "low", "backend": "CPU", "numThread": (os.cpu_count() + 1) // 2}
rt = MNN.nn.create_runtime_manager((config,))
net = MNN.nn.load_module_from_file(w, [], [], runtime_manager=rt, rearrange=True)
def torch_to_mnn(x):
return MNN.expr.const(x.data_ptr(), x.shape)
metadata = json.loads(net.get_info()["bizCode"])
# NCNN
elif ncnn:
LOGGER.info(f"Loading {w} for NCNN inference...")
check_requirements("git+https://github.com/Tencent/ncnn.git" if ARM64 else "ncnn") # requires NCNN
import ncnn as pyncnn
net = pyncnn.Net()
net.opt.use_vulkan_compute = cuda
w = Path(w)
if not w.is_file(): # if not *.param
w = next(w.glob("*.param")) # get *.param file from *_ncnn_model dir
net.load_param(str(w))
net.load_model(str(w.with_suffix(".bin")))
metadata = w.parent / "metadata.yaml"
# NVIDIA Triton Inference Server
elif triton:
check_requirements("tritonclient[all]")
from ultralytics.utils.triton import TritonRemoteModel
model = TritonRemoteModel(w)
metadata = model.metadata
# RKNN
elif rknn:
if not is_rockchip():
raise OSError("RKNN inference is only supported on Rockchip devices.")
LOGGER.info(f"Loading {w} for RKNN inference...")
check_requirements("rknn-toolkit-lite2")
from rknnlite.api import RKNNLite
w = Path(w)
if not w.is_file(): # if not *.rknn
w = next(w.rglob("*.rknn")) # get *.rknn file from *_rknn_model dir
rknn_model = RKNNLite()
rknn_model.load_rknn(w)
rknn_model.init_runtime()
metadata = Path(w).parent / "metadata.yaml"
# 可忽略-----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------↑
# 这段代码处理了当输入的模型格式不被支持时的情况。
# Any other format (unsupported)
# 检查是否为不支持的格式。如果前面的代码中没有匹配到任何支持的模型格式（如 PyTorch、ONNX、TensorRT 等），则进入这个分支。
else:
# 导入 export_formats 函数，该函数返回一个字典，包含 Ultralytics 支持的所有模型格式及其相关信息。
# def export_formats(): -> 用于生成一个字典，描述了 Ultralytics YOLO 模型支持的导出格式及其相关属性。返回一个字典，其中键为 字段名 ，值为 字段值的列表 。 -> return dict(zip(["Format", "Argument", "Suffix", "CPU", "GPU", "Arguments"], zip(*x)))
from ultralytics.engine.exporter import export_formats
# 抛出一个 TypeError 异常，提示用户输入的模型格式不被支持。
# 异常消息中包含了 ：
# 输入的模型路径 w 。
# Ultralytics 支持的模型格式列表，通过调用 export_formats()['Format'] 获取。
raise TypeError(
f"model='{w}' is not a supported model format. Ultralytics supports: {export_formats()['Format']}\n" # model='{w}' 不是受支持的模型格式。Ultralytics 支持：{export_formats()['Format']}。
f"See https://docs.ultralytics.com/modes/predict for help." # 请参阅https://docs.ultralytics.com/modes/predict 寻求帮助。
)
# 这段代码的主要作用是。检查是否为不支持的格式：如果输入的模型格式不被支持，则进入这个分支。导入 export_formats 函数：获取 Ultralytics 支持的所有模型格式。抛出异常：提示用户输入的模型格式不被支持，并列出支持的格式。通过这段代码， AutoBackend 类能够明确地告知用户，输入的模型格式不被支持，并提供支持的格式列表，帮助用户解决问题。
# 这段代码的目的是加载和处理模型的外部元数据（通常存储在 YAML 文件中），并根据元数据更新一些关键的模型属性。
# Load external metadata YAML 加载外部元数据 YAML 。
# 检查元数据文件是否存在。
# 检查 metadata 是否为字符串或 Path 对象，并且对应的文件是否存在。
if isinstance(metadata, (str, Path)) and Path(metadata).exists():
# 如果文件存在，则使用 yaml_load 函数加载 YAML 文件内容到 metadata 变量中。 yaml_load 是一个自定义函数，用于加载 YAML 文件并将其内容解析为 Python 字典。
# def yaml_load(file="data.yaml", append_filename=False): -> 用于加载和解析YAML文件。它还提供了一些额外的功能，例如清理文件内容中的特殊字符，并可以选择将文件名添加到解析后的数据中。返回解析后的YAML数据，可能包含文件名（如果 append_filename 为 True ）。 -> return data
metadata = yaml_load(metadata)
# 处理元数据。
# 如果 metadata 是一个字典。
if metadata and isinstance(metadata, dict):
# 则遍历其键值对。
for k, v in metadata.items():
# 如果键是 stride 或 batch ，则将对应的值转换为整数。
if k in {"stride", "batch"}:
metadata[k] = int(v)
# 如果键是 imgsz 、 names 、 kpt_shape 或 args ，并且值是字符串，则使用 eval 将字符串转换为实际的 Python 对象（如列表或字典）。
elif k in {"imgsz", "names", "kpt_shape", "args"} and isinstance(v, str):
metadata[k] = eval(v)
# 更新模型属性。从元数据中提取并更新模型的属性。
# 模型的步幅。
stride = metadata["stride"]
# 模型的任务类型（如检测、分割等）。
task = metadata["task"]
# 模型的批量大小。
batch = metadata["batch"]
# 模型的输入图像大小。
imgsz = metadata["imgsz"]
# 模型的类别名称。
names = metadata["names"]
# 模型的关键点形状（如果存在）。
kpt_shape = metadata.get("kpt_shape")
# 模型是否为端到端模型（通过 args 字典中的 nms 键值判断）。
end2end = metadata.get("args", {}).get("nms", False)
# 如果元数据不存在。
# 如果元数据文件不存在，并且模型不是 PyTorch、Triton 或 PyTorch 模块，则记录一个警告日志，提示用户未找到元数据。 这里假设 PyTorch、Triton 和 PyTorch 模块不需要外部元数据文件。
elif not (pt or triton or nn_module):
LOGGER.warning(f"WARNING ⚠️ Metadata not found for 'model={weights}'") # 警告⚠️未找到“model={weights}”的元数据。
# 这段代码的主要作用是。加载元数据文件：如果存在元数据文件，则加载并解析其内容。处理元数据：将元数据中的某些值转换为适当的类型（如整数或 Python 对象）。更新模型属性：根据元数据更新模型的步幅、任务类型、批量大小、输入图像大小、类别名称等属性。记录警告：如果未找到元数据文件且模型不是 PyTorch、Triton 或 PyTorch 模块，则记录警告。通过这段代码， AutoBackend 类能够根据外部元数据文件动态更新模型的属性，从而确保模型的正确性和灵活性。
# 这段代码的目的是确保模型的类别名称（ names ）是有效的，并禁用模型的梯度计算，最后将所有局部变量分配给类的属性。
# Check names
# 检查 names 是否在局部变量中。如果不存在，说明之前没有成功加载或定义类别名称。
if "names" not in locals(): # names missing
# 如果 names 缺失，则调用 default_class_names(data) 函数，使用默认的类别名称。 data 参数通常是一个路径，指向包含类别名称的 data.yaml 文件。
# def default_class_names(data=None): -> 用于从输入的 YAML 文件中提取类别名称，或者在无法提取时返回默认的类别名称。 -> return yaml_load(check_yaml(data))["names"] / return {i: f"class{i}" for i in range(999)} # return default if above errors
names = default_class_names(data)
# 调用 check_class_names(names) 函数，对类别名称进行验证和处理。 这个函数会确保类别名称是有效的，并将其转换为字典格式（如果需要）。例如，如果输入是一个列表，它会将其转换为字典，键为索引，值为类别名称。
# def check_class_names(names): -> 用于检查和处理类别名称（ names ）。返回处理后的类别名称字典。 -> return names
names = check_class_names(names)
# Disable gradients 禁用梯度计算。
# 如果模型是 PyTorch 模型（ pt 为 True ）。
if pt:
# 则遍历模型的所有参数。
for p in model.parameters():
# 并将它们的 requires_grad 属性设置为 False 。
# 是为了禁用梯度计算，因为在推理阶段不需要计算梯度，这可以节省内存并提高推理速度。
p.requires_grad = False
# 将局部变量分配给类的属性。
# 使用 self.__dict__.update(locals()) 将所有局部变量分配给类的属性。 这样可以在类的其他方法中方便地访问这些变量，而不需要显式地将它们定义为类的属性。
self.__dict__.update(locals()) # assign all variables to self
# 这段代码的主要作用是。确保类别名称存在：如果类别名称缺失，则使用默认的类别名称。验证和处理类别名称：通过 check_class_names 函数确保类别名称是有效的，并将其转换为合适的格式。禁用梯度计算：在推理阶段禁用模型的梯度计算，以节省内存并提高性能。将局部变量分配给类的属性：将所有局部变量分配给类的属性，以便在类的其他方法中使用。通过这段代码， AutoBackend 类能够确保模型的类别名称是有效的，并为推理阶段做好准备。
# AutoBackend 类的 __init__ 方法的主要功能和特点如下。动态后端选择：根据模型文件的格式，动态选择合适的推理后端。支持多种格式：支持 PyTorch、ONNX、OpenVINO、TensorRT、CoreML、TensorFlow、TFLite、PaddlePaddle、MNN、NCNN、RKNN 等多种模型格式。设备兼容性：支持 CPU 和 GPU 设备，自动检测并切换。推理优化：支持半精度推理（FP16）、动态批量大小、模型融合等优化。类别名称管理：提供默认类别名称，并支持自定义类别名称。通过 AutoBackend ，用户可以轻松地在不同平台上部署和运行 Ultralytics YOLO 模型，而无需关心底层的推理引擎细节。
# 这段代码定义了 AutoBackend 类的 forward 方法，用于在不同后端（如 PyTorch、ONNX、TensorRT 等）上运行模型推理。
# 定义了 forward 方法，接受以下参数 ：
# 1.im ：输入图像张量。
# 2.augment ：是否在推理时进行数据增强，默认为 False 。
# 3.visualize ：是否可视化输出预测，默认为 False 。
# 4.embed ：可选的特征向量/嵌入列表。
# 返回值是一个元组，包含原始输出张量和处理后的输出（如果 visualize=True ）。
def forward(self, im, augment=False, visualize=False, embed=None):
# 在 YOLOv8 MultiBackend 模型上运行推理。
"""
Runs inference on the YOLOv8 MultiBackend model.
Args:
im (torch.Tensor): The image tensor to perform inference on.
augment (bool): whether to perform data augmentation during inference, defaults to False
visualize (bool): whether to visualize the output predictions, defaults to False
embed (list, optional): A list of feature vectors/embeddings to return.
Returns:
(tuple): Tuple containing the raw output tensor, and processed output for visualization (if visualize=True)
"""
# 这段代码是 AutoBackend 类的 forward 方法的一部分，主要负责处理输入图像张量 im ，并根据模型的配置和后端类型进行推理。
# 获取输入图像张量 im 的形状，分别表示 ：
# b ：批量大小（batch size）。
# ch ：通道数（channels）。
# h ：图像高度（height）。
# w ：图像宽度（width）。
b, ch, h, w = im.shape # batch, channel, height, width
# 处理数据类型。
# 检查是否启用了半精度浮点数（FP16）推理（ self.fp16 ）。 如果启用了 FP16 且输入张量的数据类型不是 torch.float16 ，则将输入张量转换为半精度浮点数（ torch.float16 ）。 这一步可以减少内存占用并提高推理速度，但可能会牺牲一些精度。
if self.fp16 and im.dtype != torch.float16:
im = im.half() # to FP16
# 处理数据格式。
# 检查是否需要将输入张量从 PyTorch 的 BCHW 格式转换为 NHWC 格式（ self.nhwc ）。 如果需要，则调用 permute 方法重新排列张量的维度顺序，将通道维度从第一个位置移动到最后一个位置。 这一步是为了适配某些后端（如 OpenCV DNN、TensorFlow Lite 等）对输入数据格式的要求。
if self.nhwc:
im = im.permute(0, 2, 3, 1) # torch BCHW to numpy BHWC shape(1,320,192,3)
# PyTorch PyTorch 模型推理。
# 检查是否是 PyTorch 模型（ self.pt ）或 PyTorch 模块（ self.nn_module ）。
if self.pt or self.nn_module:
# 如果是，则调用模型的 forward 方法进行推理。
# im ：输入图像张量。
# augment ：是否在推理时进行数据增强。
# visualize ：是否可视化输出预测。
# embed ：可选的特征向量/嵌入列表。
# 推理结果存储在变量 y 中。
y = self.model(im, augment=augment, visualize=visualize, embed=embed)
# 这段代码的主要作用是。获取输入图像的形状：提取批量大小、通道数、高度和宽度。处理数据类型：根据配置将输入张量转换为半精度浮点数（FP16）。处理数据格式：根据需要将输入张量从 BCHW 格式转换为 NHWC 格式。进行 PyTorch 模型推理：调用 PyTorch 模型的 forward 方法进行推理，并返回结果。通过这些步骤， AutoBackend 类能够确保输入数据符合模型的要求，并在 PyTorch 模型上进行高效的推理。
# 可忽略-----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------↓
# TorchScript
elif self.jit:
y = self.model(im)
# ONNX OpenCV DNN
elif self.dnn:
im = im.cpu().numpy() # torch to numpy
self.net.setInput(im)
y = self.net.forward()
# ONNX Runtime
elif self.onnx or self.imx:
if self.dynamic:
im = im.cpu().numpy() # torch to numpy
y = self.session.run(self.output_names, {self.session.get_inputs()[0].name: im})
else:
if not self.cuda:
im = im.cpu()
self.io.bind_input(
name="images",
device_type=im.device.type,
device_id=im.device.index if im.device.type == "cuda" else 0,
element_type=np.float16 if self.fp16 else np.float32,
shape=tuple(im.shape),
buffer_ptr=im.data_ptr(),
)
self.session.run_with_iobinding(self.io)
y = self.bindings
if self.imx:
# boxes, conf, cls
y = np.concatenate([y[0], y[1][:, :, None], y[2][:, :, None]], axis=-1)
# OpenVINO
elif self.xml:
im = im.cpu().numpy() # FP32
if self.inference_mode in {"THROUGHPUT", "CUMULATIVE_THROUGHPUT"}: # optimized for larger batch-sizes
n = im.shape[0] # number of images in batch
results = [None] * n # preallocate list with None to match the number of images
def callback(request, userdata):
"""Places result in preallocated list using userdata index."""
results[userdata] = request.results
# Create AsyncInferQueue, set the callback and start asynchronous inference for each input image
async_queue = self.ov.runtime.AsyncInferQueue(self.ov_compiled_model)
async_queue.set_callback(callback)
for i in range(n):
# Start async inference with userdata=i to specify the position in results list
async_queue.start_async(inputs={self.input_name: im[i : i + 1]}, userdata=i) # keep image as BCHW
async_queue.wait_all() # wait for all inference requests to complete
y = np.concatenate([list(r.values())[0] for r in results])
else: # inference_mode = "LATENCY", optimized for fastest first result at batch-size 1
y = list(self.ov_compiled_model(im).values())
# TensorRT
elif self.engine:
if self.dynamic and im.shape != self.bindings["images"].shape:
if self.is_trt10:
self.context.set_input_shape("images", im.shape)
self.bindings["images"] = self.bindings["images"]._replace(shape=im.shape)
for name in self.output_names:
self.bindings[name].data.resize_(tuple(self.context.get_tensor_shape(name)))
else:
i = self.model.get_binding_index("images")
self.context.set_binding_shape(i, im.shape)
self.bindings["images"] = self.bindings["images"]._replace(shape=im.shape)
for name in self.output_names:
i = self.model.get_binding_index(name)
self.bindings[name].data.resize_(tuple(self.context.get_binding_shape(i)))
s = self.bindings["images"].shape
assert im.shape == s, f"input size {im.shape} {'>' if self.dynamic else 'not equal to'} max model size {s}"
self.binding_addrs["images"] = int(im.data_ptr())
self.context.execute_v2(list(self.binding_addrs.values()))
y = [self.bindings[x].data for x in sorted(self.output_names)]
# CoreML
elif self.coreml:
im = im[0].cpu().numpy()
im_pil = Image.fromarray((im * 255).astype("uint8"))
# im = im.resize((192, 320), Image.BILINEAR)
y = self.model.predict({"image": im_pil}) # coordinates are xywh normalized
if "confidence" in y:
raise TypeError(
"Ultralytics only supports inference of non-pipelined CoreML models exported with "
f"'nms=False', but 'model={w}' has an NMS pipeline created by an 'nms=True' export."
)
# TODO: CoreML NMS inference handling
# from ultralytics.utils.ops import xywh2xyxy
# box = xywh2xyxy(y['coordinates'] * [[w, h, w, h]]) # xyxy pixels
# conf, cls = y['confidence'].max(1), y['confidence'].argmax(1).astype(np.float32)
# y = np.concatenate((box, conf.reshape(-1, 1), cls.reshape(-1, 1)), 1)
y = list(y.values())
if len(y) == 2 and len(y[1].shape) != 4: # segmentation model
y = list(reversed(y)) # reversed for segmentation models (pred, proto)
# PaddlePaddle
elif self.paddle:
im = im.cpu().numpy().astype(np.float32)
self.input_handle.copy_from_cpu(im)
self.predictor.run()
y = [self.predictor.get_output_handle(x).copy_to_cpu() for x in self.output_names]
# MNN
elif self.mnn:
input_var = self.torch_to_mnn(im)
output_var = self.net.onForward([input_var])
y = [x.read() for x in output_var]
# NCNN
elif self.ncnn:
mat_in = self.pyncnn.Mat(im[0].cpu().numpy())
with self.net.create_extractor() as ex:
ex.input(self.net.input_names()[0], mat_in)
# WARNING: 'output_names' sorted as a temporary fix for https://github.com/pnnx/pnnx/issues/130
y = [np.array(ex.extract(x)[1])[None] for x in sorted(self.net.output_names())]
# NVIDIA Triton Inference Server
elif self.triton:
im = im.cpu().numpy() # torch to numpy
y = self.model(im)
# RKNN
elif self.rknn:
im = (im.cpu().numpy() * 255).astype("uint8")
im = im if isinstance(im, (list, tuple)) else [im]
y = self.rknn_model.inference(inputs=im)
# TensorFlow (SavedModel, GraphDef, Lite, Edge TPU)
else:
im = im.cpu().numpy()
if self.saved_model: # SavedModel
y = self.model(im, training=False) if self.keras else self.model(im)
if not isinstance(y, list):
y = [y]
elif self.pb: # GraphDef
y = self.frozen_func(x=self.tf.constant(im))
else: # Lite or Edge TPU
details = self.input_details[0]
is_int = details["dtype"] in {np.int8, np.int16} # is TFLite quantized int8 or int16 model
if is_int:
scale, zero_point = details["quantization"]
im = (im / scale + zero_point).astype(details["dtype"]) # de-scale
self.interpreter.set_tensor(details["index"], im)
self.interpreter.invoke()
y = []
for output in self.output_details:
x = self.interpreter.get_tensor(output["index"])
if is_int:
scale, zero_point = output["quantization"]
x = (x.astype(np.float32) - zero_point) * scale # re-scale
if x.ndim == 3: # if task is not classification, excluding masks (ndim=4) as well
# Denormalize xywh by image size. See https://github.com/ultralytics/ultralytics/pull/1695
# xywh are normalized in TFLite/EdgeTPU to mitigate quantization error of integer models
if x.shape[-1] == 6 or self.end2end: # end-to-end model
x[:, :, [0, 2]] *= w
x[:, :, [1, 3]] *= h
if self.task == "pose":
x[:, :, 6::3] *= w
x[:, :, 7::3] *= h
else:
x[:, [0, 2]] *= w
x[:, [1, 3]] *= h
if self.task == "pose":
x[:, 5::3] *= w
x[:, 6::3] *= h
y.append(x)
# TF segment fixes: export is reversed vs ONNX export and protos are transposed
if len(y) == 2: # segment with (det, proto) output order reversed
if len(y[1].shape) != 4:
y = list(reversed(y)) # should be y = (1, 116, 8400), (1, 160, 160, 32)
if y[1].shape[-1] == 6: # end-to-end model
y = [y[1]]
else:
y[1] = np.transpose(y[1], (0, 3, 1, 2)) # should be y = (1, 116, 8400), (1, 32, 160, 160)
y = [x if isinstance(x, np.ndarray) else x.numpy() for x in y]
# 可忽略-----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------↑
# 这段代码是 AutoBackend 类的 forward 方法的一部分，主要负责处理模型的输出结果，并根据需要将其转换为 PyTorch 张量。
# 这是一段被注释掉的调试代码，用于打印输出结果 y 的类型和形状。 如果 y 是一个列表或元组，打印其 类型 和 长度 ；否则打印其类型和形状。 这段代码可以帮助开发者在调试阶段了解模型输出的具体结构。
# for x in y:
# print(type(x), len(x)) if isinstance(x, (list, tuple)) else print(type(x), x.shape) # debug shapes
# 处理输出结果。
# 检查模型的输出 y 是否是一个列表或元组。
if isinstance(y, (list, tuple)):
# 如果类别名称未定义（ len(self.names) == 999 ），并且任务是分割（ self.task == "segment" ）或输出结果包含两个元素（ len(y) == 2 ）。
if len(self.names) == 999 and (self.task == "segment" or len(y) == 2): # segments and names not defined
# 则动态生成类别名称。
# 计算类别数量 nc 。
nc = y[0].shape[1] - y[1].shape[1] - 4 # y = (1, 32, 160, 160), (1, 116, 8400)
# 并生成一个字典，键为类别索引，值为格式化的类别名称（如 "class0" 、 "class1" 等）。
self.names = {i: f"class{i}" for i in range(nc)}
# 返回单个或多个输出结果。
# 如果输出结果只有一个元素（ len(y) == 1 ），则将第一个元素转换为 PyTorch 张量并返回。 如果输出结果有多个元素，则将每个元素分别转换为 PyTorch 张量，并返回一个列表。
return self.from_numpy(y[0]) if len(y) == 1 else [self.from_numpy(x) for x in y]
# 处理单个输出结果。
# 如果输出结果不是列表或元组。
else:
# 则直接将其转换为 PyTorch 张量并返回。
return self.from_numpy(y)
# 这段代码的主要作用是。调试输出：提供了一种调试方式，用于打印输出结果的类型和形状。处理输出结果：如果输出结果是列表或元组，根据任务类型和输出结构动态生成类别名称。将每个输出结果转换为 PyTorch 张量。返回结果：根据输出结果的结构，返回单个张量或张量列表。通过这些步骤， AutoBackend 类能够灵活地处理不同模型的输出结果，并将其转换为统一的 PyTorch 张量格式，方便后续的处理和使用。
# AutoBackend 类的 forward 方法是一个多功能的推理接口，能够根据不同的后端（如 PyTorch、ONNX、TensorRT、OpenVINO 等）动态选择合适的推理引擎，并处理输入数据的格式转换（如数据类型、内存布局等）。它支持多种推理模式，包括数据增强、可视化输出等，并能够根据模型的输出结构动态生成类别名称。最终，该方法将模型的输出结果统一转换为 PyTorch 张量格式，以便于后续的处理和分析。这一设计使得 AutoBackend 类能够无缝适配多种模型格式和推理需求，极大地提高了模型部署的灵活性和效率。
# 这段代码定义了 AutoBackend 类的 from_numpy 方法，用于将 NumPy 数组转换为 PyTorch 张量，并将其移动到指定的设备上。
# 定义了一个名为 from_numpy 的方法，该方法接受一个参数。
# 1.x ：需要转换的数据。
def from_numpy(self, x):
# 将 numpy 数组转换为张量。
"""
Convert a numpy array to a tensor.
Args:
x (np.ndarray): The array to be converted.
Returns:
(torch.Tensor): The converted tensor
"""
# 使用条件表达式检查 x 是否为 NumPy 数组（ np.ndarray ）。
# 如果 x 是 NumPy 数组，则使用 torch.tensor(x) 将其转换为 PyTorch 张量。
# 使用 .to(self.device) 将张量移动到指定的设备（如 CPU 或 GPU）。
# 如果 x 不是 NumPy 数组，则直接返回 x ，不做任何转换。
return torch.tensor(x).to(self.device) if isinstance(x, np.ndarray) else x
# 这段代码的主要作用是。检查输入类型：判断输入数据是否为 NumPy 数组。转换为 PyTorch 张量：如果输入是 NumPy 数组，则将其转换为 PyTorch 张量。移动到指定设备：将转换后的张量移动到指定的设备（如 CPU 或 GPU）。返回结果：返回转换后的张量或原始数据（如果输入不是 NumPy 数组）。通过这个方法， AutoBackend 类能够灵活地处理不同类型的输入数据，并确保数据在正确的设备上进行后续处理。这使得模型的推理过程更加高效和灵活。
# 这段代码定义了 AutoBackend 类的 warmup 方法，用于通过运行一次前向传播来预热模型，确保推理性能稳定。
# 定义了一个名为 warmup 的方法，接受一个参数。
# 1.imgsz ：预热时使用的输入张量的形状，默认为 (1, 3, 640, 640) ，即一个批量大小为 1、通道数为 3、高度和宽度均为 640 的张量。
def warmup(self, imgsz=(1, 3, 640, 640)):
# 通过使用虚拟输入运行一次前向传递来预热模型。
"""
Warm up the model by running one forward pass with a dummy input.
Args:
imgsz (tuple): The shape of the dummy input tensor in the format (batch_size, channels, height, width)
"""
# 导入 torchvision 模块。这里使用了 # noqa 注释，以避免某些代码检查工具（如 flake8 ）记录导入时间。
import torchvision # noqa (import here so torchvision import time not recorded in postprocess time)
# 定义了一个元组 warmup_types ，包含所有支持预热的模型类型标志（如 PyTorch、TorchScript、ONNX、TensorRT 等）。
warmup_types = self.pt, self.jit, self.onnx, self.engine, self.saved_model, self.pb, self.triton, self.nn_module
# 检查是否支持预热（即 warmup_types 中至少有一个为 True ）。 如果模型运行在 GPU 上（ self.device.type != "cpu" ）或使用 Triton 后端（ self.triton ），则进行预热。
if any(warmup_types) and (self.device.type != "cpu" or self.triton):
# 创建一个空的张量 im ，其形状由 imgsz 指定。
# 如果启用了半精度浮点数（ self.fp16 ），则将张量的数据类型设置为 torch.half ；否则设置为 torch.float 。
# 将张量移动到指定的设备（ self.device ）。
im = torch.empty(*imgsz, dtype=torch.half if self.fp16 else torch.float, device=self.device) # input
# 如果模型是 TorchScript 模型（ self.jit ），则运行两次前向传播以预热；否则运行一次。
for _ in range(2 if self.jit else 1):
# 调用 self.forward(im) 方法进行前向传播，确保模型在后续推理中能够稳定运行。
self.forward(im) # warmup
# 这段代码的主要作用是。检查是否支持预热：确保模型类型支持预热，并且运行在 GPU 上或使用 Triton 后端。创建预热输入张量：根据指定的形状和数据类型创建一个空的输入张量，并将其移动到指定设备。运行前向传播：通过运行一次或两次前向传播来预热模型，确保模型在后续推理中能够稳定运行。通过预热，可以初始化模型的各个组件，减少首次推理时的延迟，提高推理性能。
# 这段代码定义了 AutoBackend 类的静态方法 _model_type ，用于根据模型文件的路径或名称检测模型的格式。
@staticmethod
# 定义了一个静态方法 _model_type ，接受一个参数。
# 1.p ：模型文件的路径或名称，默认值为 "path/to/model.pt" 。
def _model_type(p="path/to/model.pt"):
# 获取模型文件的路径并返回模型类型。可能的类型有 pt、jit、onnx、xml、engine、coreml、saved_model、pb、tflite、edgetpu、tfjs、ncnn 或 paddle。
"""
Takes a path to a model file and returns the model type. Possibles types are pt, jit, onnx, xml, engine, coreml,
saved_model, pb, tflite, edgetpu, tfjs, ncnn or paddle.
Args:
p: path to the model file. Defaults to path/to/model.pt
Examples:
>>> model = AutoBackend(weights="path/to/model.onnx")
>>> model_type = model._model_type() # returns "onnx"
"""
# 导入 export_formats 函数，该函数返回一个字典，包含 Ultralytics 支持的所有模型格式及其相关信息。
# def export_formats(): -> 用于生成一个字典，描述了 Ultralytics YOLO 模型支持的导出格式及其相关属性。返回一个字典，其中键为 字段名 ，值为 字段值的列表 。 -> return dict(zip(["Format", "Argument", "Suffix", "CPU", "GPU", "Arguments"], zip(*x)))
from ultralytics.engine.exporter import export_formats
# 从 export_formats 函数返回的字典中提取 Suffix 键对应的值，即支持的模型文件后缀列表。
sf = export_formats()["Suffix"] # export suffixes
# 检查输入 p 是否为 URL 或字符串。
if not is_url(p) and not isinstance(p, str):
# 如果 p 不是 URL 且不是字符串，则调用 check_suffix 函数检查 p 的后缀是否在支持的后缀列表 sf 中。
# def check_suffix(file="yolo11n.pt", suffix=".pt", msg=""): -> 用于检查文件的后缀是否符合指定的允许后缀列表。如果文件的后缀不符合要求，函数会抛出一个 AssertionError 。
check_suffix(p, sf) # checks
# 使用 Path 对象提取输入路径 p 的文件名部分。
name = Path(p).name
# 遍历支持的后缀列表 sf ，检查每个后缀是否出现在文件名 name 中，生成一个布尔值列表 types 。
types = [s in name for s in sf]
# 检查文件名是否以 .mlmodel 结尾，如果是，则将 types 列表的第 6 个元素（索引为 5）设置为 True ，以支持旧版 Apple CoreML 格式。
types[5] |= name.endswith(".mlmodel") # retain support for older Apple CoreML *.mlmodel formats
# 如果文件名同时包含 .tflite 和 .edgetpu ，则将 types 列表的第 9 个元素（索引为 8）设置为 False ，以确保 .tflite 和 .edgetpu 不同时为 True 。
types[8] &= not types[9] # tflite &= not edgetpu
# 如果 types 列表中至少有一个 True ，则将 triton 设置为 False 。
if any(types):
triton = False
# 否则，解析输入路径 p 为 URL。
else:
from urllib.parse import urlsplit
# urlsplit(urlstring, scheme='', allow_fragments=True)
# urlsplit() 函数是 Python 标准库 urllib.parse 模块中的一个函数，它用于将一个 URL 分解成几个组件。
# 参数说明 ：
# urlstring ：要分解的 URL 字符串。
# scheme ：可选参数，如果提供了这个参数， urlsplit() 会尝试将 URL 按照提供的 scheme 来解析。默认为空字符串，表示不指定方案。
# allow_fragments ：可选参数，布尔值，指定是否允许 URL 中包含片段（fragment）。默认为 True 。
# 返回值 ：
# urlsplit() 函数返回一个 ParseResult 对象，该对象包含以下属性 ：
# scheme ：URL 的协议部分，例如 "http" 或 "https"。
# netloc ：网络位置部分，包括域名和端口号（如果有的话），例如 "www.example.com:80"。
# path ：URL 的路径部分，例如 "/path/to/resource"。
# params ：查询字符串之前的参数部分（对于大多数 URL 不太常见）。
# query ：查询字符串部分，例如 "key=value"。
# fragment ：URL 的片段部分（也称为锚点），例如 "#section1"。
# 如果原始 URL 字符串中不包含某个部分，对应的属性将会是空字符串。
# 这个函数是处理 URL 的常用工具，可以帮助开发者提取 URL 的各个组成部分，或者验证 URL 的格式。
url = urlsplit(p)
# 检查其是否符合 Triton 的 URL 格式（即包含主机名、路径且协议为 http 或 grpc ），如果符合，则将 triton 设置为 True 。
triton = bool(url.netloc) and bool(url.path) and url.scheme in {"http", "grpc"}
# 返回一个 包含模型格式检测结果的列表 ，列表的最后一位表示 是否为 Triton 模型 。
return types + [triton]
# 这段代码的主要作用是。检测模型格式：根据模型文件的路径或名称，检测其格式是否为 PyTorch、ONNX、TensorRT 等支持的格式。支持旧版 CoreML 格式：特别检查文件名是否以 .mlmodel 结尾，以支持旧版 Apple CoreML 格式。排除 TFLite 和 EdgeTPU 的冲突：确保 .tflite 和 .edgetpu 不同时为 True 。检测 Triton 模型：如果输入路径是一个 URL 且符合 Triton 的 URL 格式，则识别为 Triton 模型。通过这个方法， AutoBackend 类能够动态识别模型的格式，从而选择合适的后端进行推理。
# AutoBackend 类的主要功能和特点如下。动态后端选择：根据模型文件的格式，动态选择合适的推理后端。支持多种格式：支持 PyTorch、ONNX、OpenVINO、TensorRT、CoreML、TensorFlow、TFLite、PaddlePaddle、MNN、NCNN、RKNN 等多种模型格式。设备兼容性：支持 CPU 和 GPU 设备，自动检测并切换。推理优化：支持半精度推理（FP16）、动态批量大小、模型融合等优化。类别名称管理：提供默认类别名称，并支持自定义类别名称。预热功能：通过预热模型，确保推理性能稳定。通过 AutoBackend ，用户可以轻松地在不同平台上部署和运行 Ultralytics YOLO 模型，而无需关心底层的推理引擎细节。
