Guía Completa para el Entrenamiento y Despliegue de YOLOv8 con TensorRT

Arquitectura y Características de YOLOv8

YOLOv8 represenat un avance significativo en la serie de modelos SOTA (State-of-the-Art) para detección de objetos, optimizando el rendimiento respecto a su predecesor, YOLOv5. Su arquitectura se organiza fundamentalemnte en tres componentes:

  • Backbone: Basado en los principios de CSP, sustituye los módulos C3 de versiones anteriores por el módulo C2f, mejorando el flujo de gradientes.
  • Neck (Cuello): Implementa un esquema PA-FPN. Se han eliminado convoluciones tras el muestreo ascendente (upsampling) y se han integrado bloques C2f para una fusión de características más eficiente.
  • Head (Cabezal de detección): A diferencia del diseño acoplado de YOLOv5, YOLOv8 utiliza un cabezal desacoplado (Decoupled Head), separando las tareas de clasificación y regresión.

Entre las innovaciones clave se encuentran la transición hacia un diseño Anchor-free (sin anclajes predefinidos), el uso de la pérdida DFL (Distribution Focal Loss) combinada con CIoU para la regresión, y el asignador de etiquetas Task-Aligned Assigner.

Configuración del Entorno de Trabajo

Para comenzar con el entrenamiento, es necesario configurar un entorno virtual que garantice la compatibilidad entre CUDA y PyTorch. Se recomienda el uso de Conda:

# Creación del entorno
conda create -n entorno_yolov8 python=3.9 -y
conda activate entorno_yolov8

# Instalación de dependencias básicas
pip install ultralytics
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118

Preparación del Conjunto de Datos

El formato de datos sigue el estándar de YOLO, donde cada imagen tiene un archivo de texto correspondiente con las coordenadas normalizadas. A continuación, se presenta un script optimizado para convertir anotaciones XML (Pascal VOC) al formato TXT requerido:

import os
import cv2
import xml.etree.ElementTree as ET

def transformar_coordenadas(dimensiones, caja):
    escala_w = 1.0 / dimensiones[0]
    escala_h = 1.0 / dimensiones[1]
    cx = (caja[0] + caja[2]) / 2.0
    cy = (caja[1] + caja[3]) / 2.0
    ancho = caja[2] - caja[0]
    alto = caja[3] - caja[1]
    return (cx * escala_w, cy * escala_h, ancho * escala_w, alto * escala_h)

def procesar_xml_a_yolo(ruta_xml, ruta_salida, mapa_clases):
    arbol = ET.parse(ruta_xml)
    raiz = arbol.getroot()
    size = raiz.find('size')
    w = int(size.find('width').text)
    h = int(size.find('height').text)

    with open(ruta_salida, 'w') as f:
        for obj in raiz.iter('object'):
            nombre_clase = obj.find('name').text
            if nombre_clase not in mapa_clases:
                continue
            id_clase = mapa_clases[nombre_clase]
            xmlbox = obj.find('bndbox')
            coords = (float(xmlbox.find('xmin').text), float(xmlbox.find('ymin').text),
                      float(xmlbox.find('xmax').text), float(xmlbox.find('ymax').text))
            yolo_box = transformar_coordenadas((w, h), coords)
            f.write(f"{id_clase} " + " ".join([f"{a:.6f}" for a in yolo_box]) + '\n')

# Ejemplo de estructura del archivo config.yaml
# path: /ruta/datos
# train: images/train
# val: images/val
# names:
#   0: objeto_ejemplo

Entrenamiento del Modelo

Una vez preparado el archivo de configuración YAML, el entrenamiento se puede lanzar desde la terminal o mediante Python. Para una detección de clase única, se puede forzar el parámetro single_cls:

yolo task=detect mode=train model=yolov8s.pt data=configuracion.yaml epochs=100 batch=32 device=0 single_cls=True

Conversión y Optimización para Despliegue

Para maximizar el rendimiento en producción, se recomienda convertir el peso nativo de PyTorch (.pt) a formatos optimizados como ONNX y finalmente a un motor de TensorRT (.engine).

Exportación a ONNX

yolo export model=best.pt format=onnx opset=12 simplify=True

Generación del motor TensorRT

Utilizando la herramienta trtexec incluida en el SDK de TensorRT, podemos generar un modelo optimizado en FP16 para reducir la latencia:

trtexec --onnx=best.onnx --saveEngine=modelo_final.engine --fp16 --workspace=4096

Implementación de Inferencia en Python con TensorRT

El siguiente código ilustra cómo cargar el motor serializado y realizar la inferencia utilizando la API de TensorRT y PyCUDA:

import tensorrt as trt
import pycuda.driver as cuda
import pycuda.autoinit
import numpy as np
import cv2

class InferenciaYOLO:
    def __init__(self, ruta_engine):
        self.logger = trt.Logger(trt.Logger.INFO)
        with open(ruta_engine, "rb") as f:
            self.runtime = trt.Runtime(self.logger)
            self.engine = self.runtime.deserialize_cuda_engine(f.read())
        
        self.context = self.engine.create_execution_context()
        self.inputs, self.outputs, self.bindings, self.stream = self.preparar_buffers()

    def preparar_buffers(self):
        inputs, outputs, bindings = [], [], []
        stream = cuda.Stream()
        for binding in self.engine:
            shape = self.engine.get_tensor_shape(binding)
            size = trt.volume(shape)
            dtype = trt.nptype(self.engine.get_tensor_dtype(binding))
            host_mem = cuda.pagelocked_empty(size, dtype)
            cuda_mem = cuda.mem_alloc(host_mem.nbytes)
            bindings.append(int(cuda_mem))
            if self.engine.get_tensor_mode(binding) == trt.TensorIOMode.INPUT:
                inputs.append({'host': host_mem, 'device': cuda_mem})
            else:
                outputs.append({'host': host_mem, 'device': cuda_mem})
        return inputs, outputs, bindings, stream

    def ejecutar(self, imagen):
        # Preprocesamiento (Resize, BGR2RGB, Normalización)
        img_input = cv2.resize(imagen, (640, 640)).transpose(2, 0, 1).astype(np.float32) / 255.0
        np.copyto(self.inputs[0]['host'], img_input.ravel())
        
        # Transferencia a GPU
        cuda.memcpy_htod_async(self.inputs[0]['device'], self.inputs[0]['host'], self.stream)
        # Inferencia
        self.context.execute_async_v2(bindings=self.bindings, stream_handle=self.stream.handle)
        # Recuperación de resultados
        cuda.memcpy_dtoh_async(self.outputs[0]['host'], self.outputs[0]['device'], self.stream)
        self.stream.synchronize()
        
        return self.outputs[0]['host']

Implementación en C++

En entornos de alta demanda, C++ ofrece una gestión de memoria superior. Los pasos principales incluyen la inicialización del IRuntime, la creación del contexto y el post-procesamiento mediante NMS (Non-Maximum Suppression) usando OpenCV.

// Fragmento conceptual de post-procesamiento en C++
void postProceso(float* salida, int n_anchors, int n_clases, float thresh) {
    cv::Mat res = cv::Mat(n_clases + 4, n_anchors, CV_32F, salida).t();
    for (int i = 0; i < n_anchors; i++) {
        cv::Mat fila = res.row(i);
        cv::Mat puntajes = fila.colRange(4, 4 + n_clases);
        double score;
        cv::Point class_id;
        cv::minMaxLoc(puntajes, 0, &score, 0, &class_id);
        
        if (score > thresh) {
            float x = fila.at<float>(0);
            float y = fila.at<float>(1);
            float w = fila.at<float>(2);
            float h = fila.at<float>(3);
            // Convertir a coordenadas de imagen y aplicar NMS
        }
    }
}

El uso de TensorRT permite reducir el tiempo de inferencia de YOLOv8 significativamente, logrando procesar imágenes en menos de 10-15ms en hardware compatible (como las series NVIDIA RTX o Jetson), lo que lo hace ideal para aplicaciones en tiempo real.

Etiquetas: YOLOv8 TensorRT ONNX Object Detection Computer Vision

Publicado el 7-22 20:35