AWS Trainium & Inferentia 文档

YOLOS

Hugging Face's logo
加入 Hugging Face 社区

并获得增强的文档体验

开始使用

YOLOS

概述

YOLOS 模型由 Yuxin Fang、Bencheng Liao、Xinggang Wang、Jiemin Fang、Jiyang Qi、Rui Wu、Jianwei Niu、Wenyu Liu 在 You Only Look at One Sequence: Rethinking Transformer in Vision through Object Detection 中提出。YOLOS 受 DETR 启发,提出仅利用普通的 视觉 Transformer (ViT) 进行目标检测。结果表明,一个基于编码器(encoder-only)的 Transformer 也能在 COCO 上达到 42 AP,与 DETR 和 Faster R-CNN 等更复杂的框架相似。

导出到 Neuron

要将 🤗 Transformers 模型部署到 Neuron 设备上,您首先需要编译模型并将其导出为序列化格式以进行推理。以下是两种编译模型的方法,您可以选择最适合您需求的一种。这里我们以 feature-extraction 为例。

选项 1: CLI

您可以使用 Optimum 命令行界面导出模型,如下所示:

optimum-cli export neuron --model hustvl/yolos-tiny --task object-detection --batch_size 1 yolos_object_detection_neuronx/

执行 optimum-cli export neuron --help 以显示所有命令行选项及其描述。

选项 2: Python API

from optimum.neuron import NeuronModelForObjectDetection
from transformers import AutoImageProcessor


preprocessor = AutoImageProcessor.from_pretrained("hustvl/yolos-tiny")
neuron_model = NeuronModelForObjectDetection.from_pretrained("hustvl/yolos-tiny", export=True, batch_size=1)

neuron_model.save_pretrained("yolos_object_detection_neuronx")
neuron_model.push_to_hub(
    "yolos_object_detection_neuronx", repository_id="optimum/yolos-tiny-neuronx-bs1"  # Replace with your HF Hub repo id
)

NeuronYolosForObjectDetection

class optimum.neuron.NeuronYolosForObjectDetection

< >

( model: ScriptModule config: PretrainedConfig model_save_dir: str | pathlib.Path | tempfile.TemporaryDirectory | None = None model_file_name: str | None = None preprocessors: list | None = None neuron_config: NeuronDefaultConfig | None = None **kwargs )

参数

  • config (transformers.PretrainedConfig) — PretrainedConfig 是模型配置类,包含模型的所有参数。使用配置文件初始化并不会加载与模型关联的权重,只加载配置。请查看 optimum.neuron.modeling.NeuronTracedModel.from_pretrained 方法以加载模型权重。
  • model (torch.jit._script.ScriptModule) — torch.jit._script.ScriptModule 是包含由 neuron(x) 编译器编译的嵌入式 NEFF(Neuron 可执行文件格式)的 TorchScript 模块。

Neuron 模型,顶部带有目标检测头,用于 COCO 检测等任务。

此模型继承自 ~neuron.modeling.NeuronTracedModel。请查看超类文档以了解库为其所有模型实现的通用方法(例如下载或保存)。

forward

< >

( pixel_values: Tensor **kwargs )

参数

  • pixel_values (torch.Tensor | None 类型,形状为 (batch_size, num_channels, height, width),默认为 None) — 对应当前批次图像的像素值。像素值可以通过 AutoImageProcessor 从编码图像中获取。

NeuronYolosForObjectDetection 的前向传播方法,重写了 __call__ 特殊方法。仅接受编译步骤中跟踪的输入。推理过程中提供的任何额外输入都将被忽略。要包含额外输入,请在指定这些输入的情况下重新编译模型。

示例

>>> import requests
>>> from PIL import Image
>>> from optimum.neuron import NeuronYolosForObjectDetection
>>> from transformers import AutoImageProcessor

>>> url = "http://images.cocodataset.org/val2017/000000039769.jpg"
>>> image = Image.open(requests.get(url, stream=True).raw)

>>> preprocessor = AutoImageProcessor.from_pretrained("optimum/yolos-tiny-neuronx-bs1")
>>> model = NeuronYolosForObjectDetection.from_pretrained("optimum/yolos-tiny-neuronx-bs1")

>>> inputs = preprocessor(images=image, return_tensors="pt")

>>> outputs = model(**inputs)
>>> target_sizes = torch.tensor([image.size[::-1]])
>>> results = image_processor.post_process_object_detection(outputs, threshold=0.9, target_sizes=target_sizes)[0]