Introducción
En esta guía configuramos un entorno de desarrollo de IA con el NVIDIA TAO Toolkit. Mi objetivo es entrenar localmente un modelo de detección de objetos YOLOv4 en una NVIDIA RTX 4070 Ti. Utilizo Docker para mantener el sistema limpio y genero datos de entrenamiento sintéticos para no depender de descargas externas.
¿Qué es el NVIDIA TAO Toolkit?
El TAO Toolkit es un framework desarrollado por NVIDIA para simplificar y acelerar drásticamente el proceso de creación y optimización de modelos de deep learning. Fue diseñado específicamente para casos de uso en el ámbito de Computer Vision (como la detección de objetos) y la IA conversacional.
El secreto de la eficiencia de TAO radica en su arquitectura, que permite entrenar modelos complejos mediante el concepto de transfer learning.
En lugar de entrenar desde cero un modelo como YOLOv4, a un proceso que a menudo llevaría meses y requeriría recursos de cómputo masivos, el TAO Toolkit aprovecha la eficiencia de los modelos preentrenados. NVIDIA proporciona en su plataforma NGC modelos compatibles con pruning que ya han sido entrenados en conjuntos de datos generales de gran tamaño como ImageNet (base de datos visual). Estos modelos poseen un entendimiento profundo de características generales de las imágenes, como bordes, formas y texturas. El objetivo es adaptar estos conocimientos adquiridos. El modelo solo debe aprender nuevos detalles específicos, como los datos sintéticos de entrenamiento o los objetos particulares. Este proceso, conocido como fine-tuning o transfer learning, reduce drásticamente la necesidad de nuevos datos de entrenamiento y el tiempo de cómputo requerido.
Configurar controladores y Docker
Primero necesitamos el controlador de la GPU y el motor de Docker.
Usaremos la versión oficial de Docker, no la de Ubuntu, e instalaremos el NVIDIA Container Toolkit para que Docker pueda acceder a la GPU.
Comprobar controlador NVIDIA
nvidia-smi

Instalar Docker
Aunque Ubuntu ofrece un paquete de Docker llamado docker.io en sus repositorios estándar, para entornos profesionales y modernos de desarrollo de IA como la configuración del NVIDIA TAO Toolkit se recomienda instalar la edición Community de Docker (Docker CE), conocida como docker-ce, directamente desde el repositorio de Docker. La razón principal es la actualidad y compatibilidad. La versión de Ubuntu (docker.io) es mantenida por los desarrolladores de Ubuntu, suele ser más antigua y sigue el ciclo de lanzamientos de la LTS correspondiente (p. ej., Ubuntu 24.04). En cambio, la versión oficial (docker-ce) es mantenida directamente por Docker (Mirantis) e incluye las funciones y parches más recientes. Para entrenar modelos de IA como YOLOv4 es crucial que el motor de Docker esté al día para garantizar la integración y compatibilidad óptimas con el NVIDIA Container Toolkit.
# Voraussetzungen
sudo apt update
sudo apt install -y ca-certificates curl gnupg
# Docker GPG Key & Repo hinzufügen
sudo install -m 0755 -d /etc/apt/keyrings
curl -fsSL https://download.docker.com/linux/ubuntu/gpg | sudo gpg --dearmor -o /etc/apt/keyrings/docker.gpg
sudo chmod a+r /etc/apt/keyrings/docker.gpg
echo \
"deb [arch=$(dpkg --print-architecture) signed-by=/etc/apt/keyrings/docker.gpg] https://download.docker.com/linux/ubuntu \
$(. /etc/os-release && echo "$VERSION_CODENAME") stable" | \
sudo tee /etc/apt/sources.list.d/docker.list > /dev/null
# Installation
sudo apt update
sudo apt install -y docker-ce docker-ce-cli containerd.io
Instalar NVIDIA Container Toolkit
El NVIDIA Container Toolkit (anteriormente conocido como nvidia-docker) es una interfaz y un conjunto de herramientas que permiten al runtime de Docker iniciar contenedores con acceso directo a GPUs NVIDIA. Sin este toolkit, Docker no podría ver ni usar la NVIDIA RTX 4070 Ti dentro de un contenedor. El objetivo principal del NVIDIA Container Toolkit es exponer el poder de cómputo de la GPU NVIDIA a las aplicaciones que se ejecutan en un contenedor de Docker.
El toolkit asegura que el contenedor pueda ver y usar los componentes necesarios para interactuar con la GPU. Entre ellos se incluyen:
- Las NVIDIA-Treiber
- Las bibliotecas CUDA
Las bibliotecas CUDA (Compute Unified Device Architecture) son un conjunto esencial de componentes de software que permiten a los desarrolladores aprovechar la capacidad de cómputo paralelo de las GPUs NVIDIA.
# Repo hinzufügen
curl -fsSL https://nvidia.github.io/libnvidia-container/gpgkey | sudo gpg --dearmor -o /usr/share/keyrings/nvidia-container-toolkit-keyring.gpg \
&& curl -s -L https://nvidia.github.io/libnvidia-container/stable/deb/nvidia-container-toolkit.list | \
sed 's#deb https://#deb [signed-by=/usr/share/keyrings/nvidia-container-toolkit-keyring.gpg] https://#g' | \
sudo tee /etc/apt/sources.list.d/nvidia-container-toolkit.list
# Installation & Konfiguration
sudo apt update
sudo apt install -y nvidia-container-toolkit
sudo nvidia-ctk runtime configure --runtime=docker
sudo systemctl restart docker
Configurar permisos de usuario
Normalmente tendrías que anteponer sudo a tus comandos de Docker cada vez, por ejemplo al iniciar un contenedor, lo cual resulta muy incómodo con uso frecuente en un entorno de desarrollo. Para evitarlo, agregamos el usuario actual al grupo docker. Este grupo se crea durante la instalación del motor de Docker y es el único grupo no root autorizado para comunicarse con el daemon de Docker, que se ejecuta en segundo plano con privilegios de root.
sudo usermod -aG docker $USER
newgrp docker
Prueba de funcionamiento de la integración de GPU en Docker
A continuación haremos la prueba de funcionamiento tras instalar el NVIDIA Container Toolkit. El objetivo es verificar que el runtime de Docker detecta correctamente la NVIDIA RTX 4070 Ti y expone sus recursos a los contenedores. Con el flag –gpus all indicamos al daemon de Docker que mapee todas las GPUs disponibles del sistema anfitrión, una funcionalidad que aporta el Container Toolkit, al contenedor. Además, arrancamos una imagen base oficial de NVIDIA (nvidia/cuda:latest), que incluye las bibliotecas CUDA necesarias y la herramienta de diagnóstico nvidia-smi. Luego ejecutamos nvidia-smi dentro del contenedor. Si la salida muestra los detalles de tu RTX 4070 Ti, la prueba es exitosa.
docker run --rm --runtime=nvidia --gpus all ubuntu nvidia-smi
Configurar entorno TAO
TAO es un launcher en Python que controla contenedores Docker. Lo instalaremos en un entorno virtual (venv).
Crear entorno Python
# Tools installieren
sudo apt install python3-venv python3-pip
# Workspace anlegen
mkdir -p ~/tao_workspace
cd ~/tao_workspace
# Venv erstellen und aktivieren
python3 -m venv venv
source venv/bin/activate
Instalar TAO Launcher
El paquete TAO Launcher es una aplicación wrapper o herramienta de línea de comandos que simplifica la gestión y ejecución del NVIDIA TAO Toolkit. El launcher maneja automáticamente las imágenes de Docker. En lugar de escribir comandos largos y propensos a errores, permite una sintaxis muy simplificada. El launcher, entre otras cosas, hace pull de la imagen Docker correcta y con la versión adecuada desde el registro NGC. Luego inicia el contenedor con los parámetros de runtime correctos (como –gpus all y los mounts necesarios), etc.
pip install nvidia-tao setuptools
Configurar mounts (tao_mounts.json)
El término mount describe el proceso de conectar el host al contenedor Docker aislado. Los volumes son un tipo específico de mount. El trabajo con el TAO Toolkit se basa principalmente en bind mounts. Un bind mount une un directorio específico del sistema de archivos del host, por ejemplo, la carpeta del proyecto local, directamente a una ruta dentro del contenedor. Esto es esencial para acceder a los datos de entrenamiento y para guardar permanentemente los resultados del modelo. En cambio, los volumes puros son áreas de almacenamiento administradas por Docker cuyo lugar exacto en el host es irrelevante para el usuario. La necesidad del archivo tao_mounts.json surge porque el TAO Toolkit usa bind mounts. En lugar de especificar rutas complejas en cada comando, el TAO Launcher lee este JSON. El archivo define de forma consistente qué carpetas del proyecto local deben montarse en el contenedor. Automatiza la generación de los comandos correctos y garantiza que todo el entorno de trabajo esté disponible de manera coherente.
vi ~/.tao_mounts.json
{
"Mounts": [
{
"source": "/home/aaron/tao_workspace/data",
"destination": "/workspace/tao-experiments/data"
},
{
"source": "/home/aaron/tao_workspace/results",
"destination": "/workspace/tao-experiments/results"
}
],
"DockerOptions": {
"shm_size": "16G",
"ulimits": {
"memlock": -1,
"stack": 67108864
},
"user": "1000:1000",
"network": "host"
}
}
/home/aaron/tao_workspace/data & /workspace/tao-experiments/data
Aquí se almacenan los datos de entrenamiento (imágenes, anotaciones). El contenedor accede en modo lectura a esta carpeta para usar los datos en el entrenamiento del modelo YOLOv4./home/aaron/tao_workspace/results & /workspace/tao-experiments/results
Esta ruta se destina al almacenamiento de resultados. El contenedor escribe en ella modelos entrenados, puntos de control (checkpoints), registros y resultados de evaluación. Así se garantiza la persistencia de los datos.
Inicio de sesión en NVIDIA NGC
NVIDIA GPU Cloud (NGC) es una plataforma central de NVIDIA que ofrece a desarrolladores e investigadores una colección de software y recursos para deep learning, HPC y data science. NGC proporciona los modelos preentrenados necesarios para el transfer learning con TAO. Para acceder a NGC necesitas un API Key, que debes generar previamente en la web.
docker login nvcr.io
# Username: $oauthtoken <-- aquí no introduzca el nombre de usuario de NGC ;)
# Password: <Ihr API Key>
Preparación de datos
Creamos la estructura de carpetas y generamos datos de ejemplo para probar la canalización.
mkdir -p ~/tao_workspace/data/training
mkdir -p ~/tao_workspace/results/specs
Generar datos de entrenamiento (script Python)
Man podría pensar que, para reconocimiento de imágenes en producción, hay que recopilar fotos de objetos correctos y defectuosos para entrenar un modelo de IA. Esto corresponde a la “recolección de datos tradicional”, donde se graban manualmente imágenes en la línea de producción durante meses, se anotan y se usan para el entrenamiento. Este enfoque es lento y, debido a la rareza de los defectos en fábricas modernas, a menudo ineficiente.
El enfoque con el NVIDIA TAO Toolkit difiere y utiliza datos de entrenamiento sintéticos, una práctica cada vez más estándar en Industria 4.0 y control de calidad.
El script de Python sirve para generar datos sintéticos y es un paso consciente lejos de descargar imágenes de Internet. Descargar fotos, por ejemplo, de vehículos para luego distinguir entre defectuosos y correctos, presenta desventajas: los datos externos suelen ser inconsistentes en resolución, iluminación y perspectiva, y, sobre todo, carecen de anotaciones precisas y fiables, imprescindibles para un entrenamiento exitoso con TAO. Además, en la práctica industrial los defectos son muy raros, ya que las cadenas de producción tienen alta calidad. Reunir suficientes imágenes de defectos sería extremadamente costoso y lento.
La razón de no usar solo fotos reales de la planta radica en la eficiencia y reproducibilidad del aprendizaje, logradas mediante la síntesis:
El deseo de ver objetos correctos se cumple en el enfoque sintético al tomar los modelos CAD (planos digitales) de los objetos perfectos y proyectarlos en un lienzo gris en el mundo virtual. Así se conoce exactamente cómo es el objeto correcto y se puede generar en cualquier posición.
Los datos sintéticos ofrecen una gran ventaja sobre las fotos reales. Cuando un modelo renderizado desde el plano CAD aparece sobre un fondo gris, el sistema sabe al 100% dónde está la caja delimitadora (bounding box) y qué partes están cubiertas. Esta anotación perfecta acelera enormemente el entrenamiento.
Otra ventaja clave es la capacidad de simular defectos de forma controlada y aplicarlos en miles de variaciones a los modelos perfectos renderizados, algo casi imposible en producción real.
El fondo gris garantiza que el modelo aprenda a reconocer solo el objeto y sus defectos, sin distraerse por variaciones de luz, sombras o fondos borrosos de la línea de producción.
import os
from PIL import Image
# Pfade
base_path = os.path.expanduser("~/tao_workspace/data/training")
img_dir = os.path.join(base_path, "image_2")
lbl_dir = os.path.join(base_path, "label_2")
os.makedirs(img_dir, exist_ok=True)
os.makedirs(lbl_dir, exist_ok=True)
# Dummy KITTI Label (Auto in der Mitte)
kitti_label = "car 0.00 0 0.00 200.0 100.0 400.0 300.0 1.5 1.6 3.5 0.0 0.0 10.0 0.0"
# 10 Bilder erstellen
for i in range(10):
filename = f"{i:06d}"
Image.new('RGB', (1248, 384), (100, 100, 100)).save(os.path.join(img_dir, f"{filename}.jpg"))
with open(os.path.join(lbl_dir, f"{filename}.txt"), "w") as f:
f.write(kitti_label)
print("Daten generiert.")
Instalación y ejecución de la generación de datos
pip install Pillow
python3 ~/tao_workspace/data/training/generate_data.py
El comando pip install Pillow instala la biblioteca Pillow, una de las más importantes para el procesamiento de imágenes en Python. Esta biblioteca es necesaria para que el script posterior pueda generar y manipular archivos de imagen en el formato requerido. Pillow proporciona la herramienta técnica para crear el lienzo vacío de los datos sintéticos. A continuación, el comando python3 ~/tao_workspace/data/training/generate_data.py ejecuta el script de Python creado. El script procesa la lógica de generación de datos definida. Como resultado, se crean en el directorio de destino las imágenes base sintéticas y neutrales, los marcadores grises en la cantidad especificada, y se guardan en disco. Este paso completa la preparación de los archivos de imagen, que servirán como base necesaria para insertar los modelos 3D de objetos en el flujo de trabajo siguiente.
TFRecords y el archivo de especificaciones
El TAO Toolkit está diseñado para procesar datos en formato TFRecords, el formato binario optimizado de TensorFlow. Este formato es necesario porque almacena y lee los datos de entrenamiento (imágenes generadas y sus anotaciones perfectas) de manera tan eficiente que maximiza la utilización de la GPU. TFRecords permiten un procesamiento más rápido durante el entrenamiento en comparación con archivos de imagen individuales y anotaciones en CSV, algo crucial para el transfer learning acelerado por GPU en la RTX 4070 Ti. Para convertir las imágenes a este formato binario especializado se necesita un archivo de especificaciones (spec), en este caso convert_spec.txt. Spec es la abreviatura de especificación y es un archivo de configuración de texto simple que le indica al TAO Launcher y al script de conversión qué archivos copiar, cuál es el formato de origen (p. ej., imágenes y anotaciones) y qué parámetros aplicar para la conversión a TFRecords. El archivo spec actúa como la instrucción central para el proceso de conversión, garantizando una transformación reproducible y sin errores de los datos de origen al formato optimizado TFRecords.
vi ~/tao_workspace/results/specs/convert_spec.txt
kitti_config {
root_directory_path: "/workspace/tao-experiments/data/training"
image_dir_name: "image_2"
label_dir_name: "label_2"
image_extension: ".jpg"
partition_mode: "random"
num_shards: 2
val_split: 20
num_partitions: 2
}
Befehl zum Konvertieren:
tao model yolo_v4 dataset_convert \
-d /workspace/tao-experiments/results/specs/convert_spec.txt \
-o /workspace/tao-experiments/data/tfrecords/kitti_train \
-r /workspace/tao-experiments/results/
El comando realiza la conversión de los datos preparados al formato TFRecords. El resultado es la generación de varios archivos binarios optimizados en formato TFRecords. Estos archivos se guardan en el directorio de salida especificado en el sistema anfitrión, que en el contenedor corresponde a /workspace/tao-experiments/data/tfrecords/kitti_train y en el disco local al directorio montado correspondiente. Cada uno de estos archivos binarios contiene toda la información de entrenamiento en un flujo optimizado para TensorFlow, incluyendo los datos de imagen, los lienzos grises sintéticos con los objetos renderizados, y sus anotaciones, las bounding boxes y etiquetas perfectamente anotadas. Este paso concluye la fase de preparación de datos, ya que los archivos TFRecords generados son la base final de entrenamiento para el modelo YOLOv4. Este formato optimizado puede leerse de forma extremadamente rápida y eficiente por el algoritmo de entrenamiento en los siguientes pasos del flujo de trabajo TAO, asegurando una utilización máxima de la GPU.
Entrenamiento (YOLOv4)
YOLOv4, sigla de You Only Look Once versión 4, es una arquitectura avanzada para la detección de objetos en tiempo real. A diferencia de métodos anteriores, YOLOv4 realiza la clasificación del objeto y la localización de la bounding box en una sola pasada por la red neuronal. Esto permite una velocidad de procesamiento extremadamente alta, necesaria para aplicaciones industriales como la detección de objetos defectuosos en la línea de producción. Para este caso de uso, el TAO Toolkit emplea el modelo con ResNet18 como backbone, una estructura de red más pequeña y eficiente que es ideal para el transfer learning, pues ya está bien preentrenada en datos generales de imágenes. El archivo yolo_v4_train_resnet18_kitti.txt es el archivo principal de configuración o spec para el TAO Launcher y es indispensable para entrenar el modelo YOLOv4. Define todos los hiperparámetros y rutas que el contenedor necesita durante el proceso de entrenamiento. La configuración especifica la tasa de aprendizaje, el tamaño de batch, el número de epochs, el método de optimización y más. Además, indica al TAO Toolkit dónde encontrar los datos TFRecords generados en el paso anterior y en qué carpeta almacenar los resultados del entrenamiento. En suma, controla todo el proceso de transfer learning desde la entrada de datos hasta el guardado del modelo entrenado.
Crear archivo Spec
random_seed: 42
yolov4_config {
big_anchor_shape: "[(114.94, 60.67), (159.06, 114.59), (297.59, 163.17)]"
mid_anchor_shape: "[(42.99, 31.91), (79.57, 31.75), (56.80, 56.93)]"
small_anchor_shape: "[(15.60, 13.88), (30.25, 20.25), (27.03, 40.12)]"
box_matching_iou: 0.25
matching_neutral_box_iou: 0.5
arch: "resnet"
nlayers: 18
arch_conv_blocks: 2
loss_loc_weight: 0.8
loss_neg_obj_weights: 100.0
loss_class_weights: 1.0
label_smoothing: 0.0
big_grid_xy_extend: 0.05
mid_grid_xy_extend: 0.1
small_grid_xy_extend: 0.2
freeze_bn: false
force_relu: false
}
training_config {
batch_size_per_gpu: 8
num_epochs: 80
enable_qat: false
checkpoint_interval: 10
learning_rate {
soft_start_cosine_annealing_schedule {
min_learning_rate: 1e-6
max_learning_rate: 5e-4
soft_start: 0.1
}
}
regularizer {
type: L1
weight: 3e-5
}
optimizer {
adam {
epsilon: 1e-7
beta1: 0.9
beta2: 0.999
amsgrad: false
}
}
}
eval_config {
average_precision_mode: SAMPLE
batch_size: 8
matching_iou_threshold: 0.5
}
nms_config {
confidence_threshold: 0.001
clustering_iou_threshold: 0.5
top_k: 200
}
augmentation_config {
hue: 0.1
saturation: 1.5
exposure: 1.5
vertical_flip: 0
horizontal_flip: 0.5
jitter: 0.3
output_width: 1248
output_height: 384
output_channel: 3
randomize_input_shape_period: 0
mosaic_prob: 0.5
mosaic_min_ratio: 0.2
}
dataset_config {
data_sources: {
tfrecords_path: "/workspace/tao-experiments/data/tfrecords/kitti_train*"
image_directory_path: "/workspace/tao-experiments/data/training"
}
include_difficult_in_training: true
image_extension: "jpg"
target_class_mapping {
key: "car"
value: "car"
}
validation_data_sources: {
tfrecords_path: "/workspace/tao-experiments/data/tfrecords/kitti_val*"
image_directory_path: "/workspace/tao-experiments/data/training"
}
}
Renombrar archivos (arreglo para validación)
El convertidor crea archivos con nombres …fold-000…, pero el archivo spec busca kitti_val. Los renombramos:
cd ~/tao_workspace/data/tfrecords
for f in *fold-000*; do mv "$f" "${f/kitti_train/kitti_val}"; done
cd ~/tao_workspace
Inicio del aprendizaje por transferencia
Una vez realizadas todas las preparaciones, ejecutamos el comando que inicia el transfer learning en la GPU NVIDIA. Este comando es el punto culminante del flujo de trabajo y ordena al TAO Launcher iniciar la imagen Docker de YOLOv4 preentrenada. Se proporciona el archivo spec yolo_v4_train_resnet18_kitti.txt como configuración central. El launcher monta los datos TFRecords de entrenamiento y la carpeta de resultados en el contenedor. Luego inicia el entrenamiento: el backbone ResNet18 preentrenado de YOLOv4 carga sus pesos y los ajusta según los datos sintéticos optimizados en TFRecords. Toda la carga de cómputo recae en la RTX 4070 Ti, mientras el NVIDIA Container Toolkit garantiza la comunicación eficiente entre el framework de deep learning y la GPU. Durante el entrenamiento se generan checkpoints periódicamente en la carpeta de resultados montada, que reflejan el progreso de aprendizaje del modelo. Estos checkpoints servirán luego para evaluar y exportar el modelo final. Con este comando comienza la fase en que el modelo aprende a distinguir defectos de objetos correctos de forma fiable.
tao model yolo_v4 train \
-e /workspace/tao-experiments/results/specs/yolo_v4_train_resnet18_kitti.txt \
-r /workspace/tao-experiments/results/train_run_1 \
-k my_key \
--gpus 1
Salida esperada:
Epoch 1/80
1/1 [==============================] - 21s 21s/step - loss: 2165133.0000
...
Epoch 80/80
1/1 [==============================] - 0s 215ms/step - loss: 349471.9688
Execution status: PASS
Exportación del modelo al formato ONNX
Tras completar con éxito el transfer learning, el modelo YOLOv4 entrenado debe convertirse a un formato adecuado para producción. El comando export realiza exactamente esto: indica al TAO Launcher que cargue el modelo generado y lo convierta al formato ONNX. ONNX (Open Neural Network Exchange) es un estándar abierto que permite intercambiar modelos de deep learning entre distintos frameworks y entornos de ejecución. Al exportar a ONNX, el modelo queda independiente del framework y puede cargarse en motores como NVIDIA TensorRT. TensorRT optimiza adicionalmente el modelo para lograr el máximo rendimiento y la menor latencia en GPUs NVIDIA. Este paso es crucial para aplicaciones en tiempo real, como NVIDIA DeepStream, muy utilizado en análisis de vídeo y procesamiento de imágenes en producción. El resultado del comando export es un archivo final en formato ONNX listo para integrarse en el sistema de destino, marcando el inicio de la fase de despliegue del modelo.
tao model yolo_v4 export \
-m /workspace/tao-experiments/results/train_run_1/weights/yolov4_resnet18_epoch_080.hdf5 \
-k my_key \
-o /workspace/tao-experiments/results/yolo_v4_resnet18_epoch_080.etlt \
-e /workspace/tao-experiments/results/specs/yolo_v4_train_resnet18_kitti.txt
Una vez configurada con éxito toda la canalización de entrenamiento local de IA, el usuario puede iniciar sus propios proyectos. Solo es necesario reemplazar las imágenes sintéticas por fotos propias. Para anotar objetos en estas fotos se recomienda usar una herramienta externa como labelImg para las anotaciones.
