Предобработка с ускорением на GPU с использованием NVIDIA DALI#

При развёртывании моделей Ultralytics YOLO в рабочей среде предобработка часто становится узким местом. Хотя TensorRT может выполнять вывод модели всего за несколько миллисекунд, предобработка на CPU (изменение размера, добавление отступов, нормализация) может занимать 2–10 мс на изображение, особенно при высоком разрешении. NVIDIA DALI (библиотека загрузки данных) решает эту проблему, перенося весь конвейер предобработки на GPU.

В этом руководстве показано, как создавать конвейеры DALI, в точности воспроизводящие предобработку Ultralytics YOLO, интегрировать их с model.predict(), обрабатывать видеопотоки и выполнять сквозное развёртывание с помощью Triton Inference Server.

Зачем использовать DALI для предобработки YOLO#

В типичном конвейере вывода YOLO этапы предобработки выполняются на CPU:

1. Декодирование изображения (JPEG/PNG)
2. Изменение размера с сохранением соотношения сторон
3. Добавление отступов до целевого размера (letterbox)
4. Нормализация значений пикселей с [0, 255] до [0, 1]
5. Преобразование формата из HWC в CHW

С DALI все эти операции выполняются на GPU, что устраняет узкое место на CPU. Это особенно полезно, когда:

Предварительные требования#

Установи необходимые пакеты:

pip install ultralytics
pip install --extra-index-url https://pypi.nvidia.com nvidia-dali-cuda130

Требования:

• GPU NVIDIA (вычислительная способность 5.0+ / Maxwell или новее)
• CUDA 11.0+, 12.0+ или 13.0+
• Python 3.10–3.14
• Операционная система Linux

Как работает предобработка YOLO#

Перед созданием конвейера DALI стоит разобраться, что именно Ultralytics выполняет во время предобработки. Ключевой класс — LetterBox в ultralytics/data/augment.py:

from ultralytics.data.augment import LetterBox

letterbox = LetterBox(
    new_shape=(640, 640),  # Target size
    center=True,  # Center the image (pad equally on both sides)
    stride=32,  # Stride alignment
    padding_value=114,  # Gray padding (114, 114, 114)
)

Полный конвейер предобработки в ultralytics/engine/predictor.py выполняет следующие шаги:

Операция letterbox сохраняет соотношение сторон следующим образом:

1. Вычисление масштаба: r = min(target_h / h, target_w / w)
2. Изменение размера до (round(w * r), round(h * r))
3. Заполнение оставшегося пространства серым цветом (114) до достижения целевого размера
4. Центрирование изображения, чтобы отступы равномерно распределялись с обеих сторон

Конвейер DALI для YOLO#

Рекомендуемый конвейер DALI воспроизводит стандартное поведение Ultralytics LetterBox(center=True), используемое при стандартном выводе YOLO.

Центрированный конвейер (рекомендуется, соответствует Ultralytics LetterBox)#

Эта версия в точности воспроизводит стандартную предобработку Ultralytics с центрированными отступами, соответствующими LetterBox(center=True):

from nvidia import dali
from nvidia.dali import fn, types

@dali.pipeline_def(batch_size=8, num_threads=4, device_id=0)
def yolo_dali_pipeline_centered(image_dir, target_size=640):
    """DALI pipeline replicating YOLO preprocessing with centered padding.

    Matches Ultralytics LetterBox(center=True) behavior exactly.
    """
    # Read and decode images on GPU
    jpegs, _ = fn.readers.file(file_root=image_dir, random_shuffle=False, name="Reader")
    images = fn.decoders.image(jpegs, device="mixed", output_type=types.RGB)

    # Aspect-ratio-preserving resize
    resized = fn.resize(
        images,
        resize_x=target_size,
        resize_y=target_size,
        mode="not_larger",
        interp_type=types.INTERP_LINEAR,
        antialias=False,  # Match cv2.INTER_LINEAR (no antialiasing)
    )

    # Centered padding using fn.crop with out_of_bounds_policy
    # When crop size > image size, fn.crop centers the image and pads symmetrically
    padded = fn.crop(
        resized,
        crop=(target_size, target_size),
        out_of_bounds_policy="pad",
        fill_values=114,  # YOLO padding value
    )

    # Normalize and convert layout
    output = fn.crop_mirror_normalize(
        padded,
        dtype=types.FLOAT,
        output_layout="CHW",
        mean=[0.0, 0.0, 0.0],
        std=[255.0, 255.0, 255.0],
    )
    return output

Запуск конвейера#

# Build and run the pipeline
pipe = yolo_dali_pipeline_centered(image_dir="/path/to/images", target_size=640)
pipe.build()

# Get a batch of preprocessed images
(output,) = pipe.run()

# Convert to numpy or PyTorch tensors
batch_np = output.as_cpu().as_array()  # Shape: (batch_size, 3, 640, 640)
print(f"Output shape: {batch_np.shape}, dtype: {batch_np.dtype}")
print(f"Value range: [{batch_np.min():.4f}, {batch_np.max():.4f}]")

Использование DALI с Ultralytics Predict#

Ты можешь напрямую передать предварительно обработанный тензор PyTorch в model.predict(). Если передан torch.Tensor, Ultralytics пропускает предобработку изображения (letterbox, BGR→RGB, HWC→CHW и нормализацию /255) и перед отправкой в модель выполняет только перенос на устройство и приведение типа данных.

Поскольку в этом случае Ultralytics не имеет доступа к исходным размерам изображения, координаты рамок обнаруженных объектов возвращаются в пространстве letterbox размером 640×640. Чтобы преобразовать их обратно в координаты исходного изображения, используй scale_boxes, который применяет ту же логику округления, что и LetterBox:

from ultralytics.utils.ops import scale_boxes

# boxes: tensor of shape (N, 4) in xyxy format, in 640x640 letterboxed coords
# Scale boxes from letterboxed (640, 640) back to original (orig_h, orig_w)
boxes = scale_boxes((640, 640), boxes, (orig_h, orig_w))

Это относится ко всем внешним путям предобработки — прямому вводу тензора, видеопотокам и развёртыванию Triton.

from nvidia.dali.plugin.pytorch import DALIGenericIterator

from ultralytics import YOLO

# Load model
model = YOLO("yolo26n.pt")

# Create DALI iterator
pipe = yolo_dali_pipeline_centered(image_dir="/path/to/images", target_size=640)
pipe.build()
dali_iter = DALIGenericIterator(pipe, ["images"], reader_name="Reader")

# Run inference with DALI-preprocessed tensors
for batch in dali_iter:
    images = batch[0]["images"]  # Already on GPU, shape (B, 3, 640, 640)
    results = model.predict(images, verbose=False)
    for result in results:
        print(f"Detected {len(result.boxes)} objects")

DALI с видеопотоками#

Для обработки видео в реальном времени используй fn.external_source, чтобы передавать кадры из любого источника — OpenCV, GStreamer или пользовательских библиотек захвата:

from nvidia import dali
from nvidia.dali import fn, types

@dali.pipeline_def(batch_size=1, num_threads=4, device_id=0)
def yolo_video_pipeline(target_size=640):
    """DALI pipeline for processing video frames from external source."""
    # External source for feeding frames from OpenCV, GStreamer, etc.
    frames = fn.external_source(device="cpu", name="input")
    frames = fn.reshape(frames, layout="HWC")

    # Move to GPU and preprocess
    frames_gpu = frames.gpu()
    resized = fn.resize(
        frames_gpu,
        resize_x=target_size,
        resize_y=target_size,
        mode="not_larger",
        interp_type=types.INTERP_LINEAR,
        antialias=False,
    )
    padded = fn.crop(
        resized,
        crop=(target_size, target_size),
        out_of_bounds_policy="pad",
        fill_values=114,
    )
    output = fn.crop_mirror_normalize(
        padded,
        dtype=types.FLOAT,
        output_layout="CHW",
        mean=[0.0, 0.0, 0.0],
        std=[255.0, 255.0, 255.0],
    )
    return output

Triton Inference Server с DALI#

Для развёртывания в рабочей среде объедини предобработку DALI с выводом TensorRT в Triton Inference Server, используя ансамблевую модель. Это полностью устраняет предобработку на CPU: на вход поступают необработанные байты JPEG, а на выходе получаются обнаружения; всё обрабатывается на GPU.

Структура репозитория моделей#

model_repository/
├── dali_preprocessing/
│   ├── 1/
│   │   └── model.dali
│   └── config.pbtxt
├── yolo_trt/
│   ├── 1/
│   │   └── model.plan
│   └── config.pbtxt
└── ensemble_dali_yolo/
    ├── 1/                  # Empty directory (required by Triton)
    └── config.pbtxt

Шаг 1: Создание конвейера DALI#

Сериализуй конвейер DALI для бэкенда Triton DALI:

from nvidia import dali
from nvidia.dali import fn, types

@dali.pipeline_def(batch_size=8, num_threads=4, device_id=0)
def triton_dali_pipeline():
    """DALI preprocessing pipeline for Triton deployment."""
    # Input: raw encoded image bytes from Triton
    images = fn.external_source(device="cpu", name="DALI_INPUT_0")
    images = fn.decoders.image(images, device="mixed", output_type=types.RGB)

    resized = fn.resize(
        images,
        resize_x=640,
        resize_y=640,
        mode="not_larger",
        interp_type=types.INTERP_LINEAR,
        antialias=False,
    )
    padded = fn.crop(
        resized,
        crop=(640, 640),
        out_of_bounds_policy="pad",
        fill_values=114,
    )
    output = fn.crop_mirror_normalize(
        padded,
        dtype=types.FLOAT,
        output_layout="CHW",
        mean=[0.0, 0.0, 0.0],
        std=[255.0, 255.0, 255.0],
    )
    return output

# Serialize pipeline to model repository
pipe = triton_dali_pipeline()
pipe.serialize(filename="model_repository/dali_preprocessing/1/model.dali")

Шаг 2: Экспорт YOLO в TensorRT#

from ultralytics import YOLO

model = YOLO("yolo26n.pt")
model.export(
    format="engine", imgsz=640, quantize=16, batch=8, dynamic=True, nms=False
)  # NMS-free (N, 300, 6); TensorRT >= 8.5
# Copy the .engine file to model_repository/yolo_trt/1/model.plan

Шаг 3: Настройка Triton#

dali_preprocessing/config.pbtxt:

name: "dali_preprocessing"
backend: "dali"
max_batch_size: 8
input [
  {
    name: "DALI_INPUT_0"
    data_type: TYPE_UINT8
    dims: [ -1 ]
  }
]
output [
  {
    name: "DALI_OUTPUT_0"
    data_type: TYPE_FP32
    dims: [ 3, 640, 640 ]
  }
]

yolo_trt/config.pbtxt:

name: "yolo_trt"
platform: "tensorrt_plan"
max_batch_size: 8
input [
  {
    name: "images"
    data_type: TYPE_FP32
    dims: [ 3, 640, 640 ]
  }
]
output [
  {
    name: "output0"
    data_type: TYPE_FP32
    dims: [ 300, 6 ]
  }
]

ensemble_dali_yolo/config.pbtxt:

name: "ensemble_dali_yolo"
platform: "ensemble"
max_batch_size: 8
input [
  {
    name: "INPUT"
    data_type: TYPE_UINT8
    dims: [ -1 ]
  }
]
output [
  {
    name: "OUTPUT"
    data_type: TYPE_FP32
    dims: [ 300, 6 ]
  }
]
ensemble_scheduling {
  step [
    {
      model_name: "dali_preprocessing"
      model_version: -1
      input_map {
        key: "DALI_INPUT_0"
        value: "INPUT"
      }
      output_map {
        key: "DALI_OUTPUT_0"
        value: "preprocessed_image"
      }
    },
    {
      model_name: "yolo_trt"
      model_version: -1
      input_map {
        key: "images"
        value: "preprocessed_image"
      }
      output_map {
        key: "output0"
        value: "OUTPUT"
      }
    }
  ]
}

Шаг 4: Отправка запросов на вывод#

import numpy as np
import tritonclient.http as httpclient

client = httpclient.InferenceServerClient(url="localhost:8000")

# Load image as raw bytes (JPEG/PNG encoded)
image_data = np.fromfile("image.jpg", dtype="uint8")
image_data = np.expand_dims(image_data, axis=0)  # Add batch dimension

# Create input
input_tensor = httpclient.InferInput("INPUT", image_data.shape, "UINT8")
input_tensor.set_data_from_numpy(image_data)

# Run inference through the ensemble
result = client.infer(model_name="ensemble_dali_yolo", inputs=[input_tensor])
detections = result.as_numpy("OUTPUT")  # Shape: (1, 300, 6) -> [x1, y1, x2, y2, conf, class_id]

# Filter by confidence (no NMS needed for the nms=False export)
detections = detections[0]  # First image
detections = detections[detections[:, 4] > 0.25]  # Confidence threshold
print(f"Detected {len(detections)} objects")

Поддерживаемые задачи#

Предобработка DALI работает со всеми задачами YOLO, использующими стандартный конвейер LetterBox:

Ограничения#

• Только Linux: DALI не поддерживает Windows или macOS
• Требуется NVIDIA GPU: резервного варианта только для CPU нет
• Статический конвейер: структура конвейера определяется во время сборки и не может изменяться динамически
• fn.pad работает только вправо/вниз: используй fn.crop с out_of_bounds_policy="pad" для центрированного заполнения
• Нет режима rect: конвейеры DALI создают выходные данные фиксированного размера (например, 640×640). Режим rect auto=True, создающий выходные данные переменного размера (например, 384×640), не поддерживается. Обрати внимание: хотя TensorRT поддерживает динамические размеры входных данных, конвейер DALI фиксированного размера естественным образом сочетается с движком фиксированного размера, обеспечивая максимальную пропускную способность
• Память при использовании нескольких экземпляров: использование instance_group с count > 1 в Triton может привести к высокому потреблению памяти. Используй группу экземпляров по умолчанию для модели DALI

Часто задаваемые вопросы#