Einleitung
In diesem Guide richten wir eine KI-Entwicklungsumgebung mit dem NVIDIA TAO Toolkit ein. Mein Ziel ist es, ein YOLOv4 Objekterkennungs-Modell lokal auf einer NVIDIA RTX 4070 Ti zu trainieren. Ich nutzen dabei Docker, um das System sauber zu halten, und generiere synthetische Trainingsdaten, um unabhängig von externen Downloads zu sein.
Was ist das NVIDIA TAO Toolkit?
Das TAO Toolkit ist ein Framework, das von NVIDIA entwickelt wurde, um den Prozess der Erstellung und Optimierung von Deep-Learning-Modellen drastisch zu vereinfachen und zu beschleunigen. Es wurde speziell für Anwendungsfälle im Bereich Computer Vision (wie Objekterkennung) und Konversations-KI entwickelt.
Der Schlüssel zur Effizienz von TAO liegt in seiner Architektur, die es ermöglicht, komplexe Modelle mithilfe des Konzepts des Transfer Learnings zu trainieren.
Anstatt ein komplexes Modell wie YOLOv4 von Grund auf neu zu trainieren ein Prozess der oft Monate dauern und massive Rechenressourcen erfordern würde nutzt das TAO Toolkit die Effizienz von Pre trained Models. NVIDIA stellt über seine NGC-Plattform sogenannte Pruning fähige Modelle zur Verfügung die bereits auf riesigen allgemeinen Datensätzen wie ImageNet (visuelle DB) vortrainiert wurden. Diese Modelle besitzen dadurch ein tiefes Verständnis für allgemeine Bildmerkmale wie Kanten Formen und Texturen. Das Ziel ist es diese bereits erworbenen Kenntnisse zu adaptieren. Das Modell muss nun lediglich neue spezifische Details lernen die die einzigartigen Daten beispielsweise synthetische Trainingsdaten oder spezifische Objekte betreffen. Dieser Prozess der als Feintuning oder Transfer Learning bezeichnet wird reduziert den Bedarf an neuen Trainingsdaten und die benötigte Rechenzeit dramatisch.
Treiber & Docker einrichten
Zuerst benötigen wir den Grafiktreiber und die Docker-Engine.
Wir nutzen die offizielle Docker-Version, nicht die von Ubuntu, und installieren das NVIDIA Container Toolkit, damit Docker auf die GPU zugreifen kann.
NVIDIA Treiber prüfen
nvidia-smi

Docker installieren
Obwohl Ubuntu ein Docker-Paket namens docker.io in seinen Standard-Repositories anbietet, wird für professionelle und moderne KI-Entwicklungsumgebungen wie die Einrichtung des NVIDIA TAO Toolkits die Installation der offiziellen Docker Community Edition (Docker CE), bekannt als docker-ce, direkt vom Docker-Repository empfohlen. Der Hauptgrund liegt in der Aktualität und Kompatibilität. Die Ubuntu-Version (docker.io) wird von den Ubuntu-Entwicklern gewartet, ist oft älter und hält sich an den Veröffentlichungszyklus der jeweiligen LTS-Version (z.B. Ubuntu 24.04). Im Gegensatz dazu wird die offizielle Version (docker-ce) direkt von Docker (Mirantis) gewartet und enthält die neuesten stabilen Funktionen und Patches. Für das Training von KI-Modellen wie YOLOv4 ist es entscheidend, dass die Docker-Engine auf dem neuesten Stand ist, um die optimale Integration und Kompatibilität mit dem NVIDIA Container Toolkit zu gewährleisten.
# 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
NVIDIA Container Toolkit installieren
Das NVIDIA Container Toolkit (früher bekannt als nvidia-docker) ist eine Schnittstelle und eine Reihe von Tools, die es der Docker-Laufzeitumgebung (dem „Runtime“) ermöglichen, Container zu starten, die direkt auf NVIDIA GPUs zugreifen können. Ohne dieses Toolkit würde Docker die NVIDIA RTX 4070 Ti in einem Container nicht sehen oder nutzen können. Das Hauptziel des NVIDIA Container Toolkits ist es, die Rechenleistung der NVIDIA-Grafikkarte für Anwendungen bereitzustellen, die in einem Docker-Container laufen.
Das Toolkit stellt sicher, dass der Container die notwendigen Komponenten sehen und nutzen kann, um mit der GPU zu interagieren. Dazu gehören:
- Die NVIDIA-Treiber
- Die CUDA-Bibliotheken
Die CUDA-Bibliotheken (Compute Unified Device Architecture) sind eine essenzielle Sammlung von Software-Komponenten, die es Entwicklern ermöglichen, die parallele Rechenleistung von NVIDIA GPUs zu nutzen.
# 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
Benutzerrechte setzen
Normalerweise müssten Sie jedes Mal sudo vor Ihre Docker-Befehle setzen, beispielsweise um einen Container zu starten, was bei häufiger Nutzung in einer Entwicklungsumgebung sehr umständlich ist. Um dies zu vermeiden, weisen wir dem aktuellen Benutzer die Mitgliedschaft in der docker-Gruppe zu. Diese Gruppe wird während der Installation der Docker Engine erstellt und ist die einzige nicht-Root-Gruppe, die berechtigt ist, mit dem Docker-Daemon zu kommunizieren, der im Hintergrund mit Root-Rechten läuft.
sudo usermod -aG docker $USER
newgrp docker
Funktionstest der GPU-Integration in Docker
Als Nächstes wird der Funktionstest nach der Installation des NVIDIA Container Toolkits durchgeführt. Das Ziel ist die Verifizierung, dass die Docker-Laufzeitumgebung die NVIDIA RTX 4070 Ti erfolgreich erkennen und deren Ressourcen für Container bereitstellen kann. Durch den Flag –gpus all wird der Docker-Daemon angewiesen, alle verfügbaren GPUs des Host-Systems, eine Funktionalität, die das Container Toolkit hinzufügt, in den Container zu mappen. Gleichzeitig wird ein offizielles NVIDIA-Basis-Image (nvidia/cuda:latest), das die notwendigen CUDA-Bibliotheken und das Diagnosetool nvidia-smi enthält, gestartet. Der Befehl nvidia-smi wird anschließend innerhalb des Containers ausgeführt. Wenn die Ausgabe dieses Befehls die Details Ihrer RTX 4070 Ti anzeigt, ist der Test erfolgreich.
docker run --rm --runtime=nvidia --gpus all ubuntu nvidia-smi
TAO Umgebung einrichten
TAO ist ein Python-Launcher, der Docker-Container steuert. Wir installieren ihn in einer virtuellen Umgebung (Venv).
Python Umgebung erstellen
# 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
TAO Launcher installieren
Das TAO Launcher Paket ist eine Wrapper-Anwendung oder ein Befehlszeilen-Tool, das die Verwaltung und Ausführung des eigentlichen NVIDIA TAO Toolkits vereinfacht. Der Launcher ist für die automatische Handhabung der Docker-Images zuständig. Anstatt lange, fehleranfällige Docker-Befehle eingeben zu müssen, ermöglicht der Launcher eine stark vereinfachte Syntax. Der Launcher zieht (pullt) unter anderem automatisch das korrekte, versionsspezifische Docker-Image von der NGC Registry. Er startet den Docker-Container mit den korrekten Laufzeitparametern (wie –gpus all und den notwendigen Volume Mounts) usw.
pip install nvidia-tao setuptools
Mounts konfigurieren (tao_mounts.json)
Der Begriff Mount beschreibt allgemein den Vorgang eine Verbindung zwischen dem Host System und dem isolierten Docker Container herzustellen. Es handelt sich um die Methode. Volumes stellen eine spezifische Art von Mount dar. Die Arbeit mit dem TAO Toolkit setzt primär auf Bind Mounts. Ein Bind Mount verbindet ein spezifisches Verzeichnis auf dem Host Dateisystem zum Beispiel den lokalen Projektordner direkt mit einem Pfad im Container. Dies ist unerlässlich da so auf die Trainingsdaten zugegriffen und die Modellergebnisse dauerhaft gespeichert werden. Im Gegensatz dazu sind reine Volumes von Docker verwaltete Speicherbereiche deren genauer Speicherort auf dem Host für den Benutzer irrelevant ist. Die Notwendigkeit der tao mounts json Datei ergibt sich daraus dass das TAO Toolkit Bind Mounts verwendet. Anstatt bei jedem Befehl manuelle und komplexe Pfadangaben machen zu müssen liest der TAO Launcher diese JSON Datei. Die Datei definiert konsistent welche lokalen Projektordner als Bind Mounts in den Container eingebunden werden sollen. Sie automatisiert somit die Generierung der korrekten Befehle und stellt sicher dass die gesamte Arbeitsumgebung konsistent verfügbar ist.
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
Hier werden die Trainingsdaten (Bilder, Annotationen) gespeichert. Der Container greift lesend auf diesen Ordner zu, um die Daten für das Training des YOLOv4-Modells zu verwenden. - /home/aaron/tao_workspace/results & /workspace/tao-experiments/results
Dieser Pfad dient der Ergebnisspeicherung. Der Container schreibt die trainierten Modelle, Checkpoints, Logs und Evaluationsergebnisse auf diesen Ordner des Host-Systems. Dies gewährleistet die Datenpersistenz.
Login bei NVIDIA NGC
Die NVIDIA GPU Cloud (NGC) ist eine zentrale Plattform von NVIDIA, die Entwicklern und Forschern eine umfassende Sammlung von Software und Ressourcen für das Deep Learning, High Performance Computing (HPC) und die Datenwissenschaft zur Verfügung stellt. NGC stellt die Pre trained Models (vortrainierte Modelle) bereit, die für das Transfer Learning mit TAO benötigt werden. Für den Zugriff auf NGC wird ein API-Key benötigt, welches wir zuvor auf der Webseite erstellen müssen.
docker login nvcr.io
# Username: $oauthtoken <-- hier nicht den NGC-Username eingeben ;)
# Password: <Ihr API Key>
Datenvorbereitung
Wir erstellen eine Ordnerstruktur und generieren Dummy-Daten, um die Pipeline zu testen.
mkdir -p ~/tao_workspace/data/training
mkdir -p ~/tao_workspace/results/specs
Trainingsdaten generieren (Python Skript)
Man könnte annehmen, dass man bei der Bilderkennung in der Produktion Bilder von korrekten und defekten Objekten sammeln muss, um ein KI-Modell zu trainieren. Dies entspricht der herkömmlichen Datensammlung („Traditional Data Collection“), bei der man monatelang Bilder direkt am Fließband aufzeichnet, manuell annotiert und diese für das Training verwendet. Dieser herkömmliche Ansatz ist zeitaufwändig und aufgrund der Seltenheit von Produktionsfehlern in modernen Fabriken oft ineffizient.
Doch der Ansatz mit dem NVIDIA TAO Toolkit weicht davon ab und nutzt synthetische Trainingsdaten, ein Vorgehen, das in der Industrie 4.0 und der Qualitätskontrolle zunehmend zum Standard wird.
Das erstellte Python Skript dient zur Generierung von synthetischen Trainingsdaten und ist ein bewusster Schritt weg vom direkten Download von Bildern aus dem Internet. Der Grund dafür ist, dass das Herunterladen von Bildern, beispielsweise von Fahrzeugen, die später in defekte und korrekte unterschieden werden sollen, erhebliche Nachteile mit sich bringt. Solche externen Daten sind oft inkonsistent in Auflösung, Beleuchtung und Perspektive, und vor allem fehlt ihnen die präzise und zuverlässige Annotation der Objekte, die für ein erfolgreiches Training mit dem NVIDIA TAO Toolkit unerlässlich ist. Das größte Problem in der industriellen Praxis ist zudem, dass Fehlerbilder extrem selten sind, da die meisten Fließbänder eine hohe Produktqualität aufweisen. Es wäre extrem zeitaufwändig und kostspielig, genügend Defekte für ein robustes Modelltraining zu sammeln.
Der Grund, warum man nicht einfach nur reale Bilder von der Fabrikhalle verwendet, liegt in der Effizienz und Reproduzierbarkeit des Lernprozesses, was durch die Synthese erreicht wird:
- Der Wunsch, die korrekten Objekte zu sehen, wird im synthetischen Ansatz erfüllt, indem die CAD Modelle (digitale Baupläne) der perfekten Objekte genommen und diese in der virtuellen Welt auf eine graue Leinwand projiziert werden. Dadurch ist exakt bekannt, wie das korrekte Objekt aussieht, und es kann in jeder gewünschten Position generiert werden.
- Synthetische Daten bieten einen massiven Vorteil gegenüber realen Fotos. Wenn ein Modell aus dem CAD-Plan auf ein graues Bild gerendert wird, weiß das System zu 100% genau, wo die Bounding Box sitzt und welche Teile davon bedeckt sind. Diese perfekte Annotation beschleunigt das Training ungemein.
- Ein weiterer entscheidender Vorteil ist die Fähigkeit, den Defekt gezielt zu simulieren und in Tausenden von Variationen auf die gerenderten perfekten Modelle anzuwenden, was in der realen Produktion kaum möglich ist.
- Die graue Leinwand stellt sicher, dass das Modell lernt, nur das Objekt selbst und seine Defekte zu erkennen, und nicht durch wechselnde Lichtverhältnisse, Schatten oder verschwommene Hintergründe an der Produktionslinie abgelenkt wird.
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.")
Installation und Ausführung der Datengenerierung
pip install Pillow
python3 ~/tao_workspace/data/training/generate_data.py
Der Befehl pip install Pillow dient der Installation der Pillow Bibliothek, welche eine der wichtigsten Bildverarbeitungsbibliotheken für Python darstellt. Diese Bibliothek ist notwendig, damit das nachfolgende Skript Bilddateien im erforderlichen Format generieren und manipulieren kann. Pillow stellt somit das technische Werkzeug bereit, um die leere Leinwand der synthetischen Daten zu erzeugen. Anschließend startet der Befehl python3 ~/tao_workspace/data/training/generate_data.py die eigentliche Ausführung des zuvor erstellten Python Skripts. Das Skript beginnt nun, die darin definierte Logik zur Datenerzeugung abzuarbeiten. Als Ergebnis werden im Zielverzeichnis die synthetischen, neutralen Basisbilder die grauen Platzhalter in der definierten Anzahl erstellt und auf der Festplatte gespeichert. Dieser Schritt schließt die Vorbereitung der Bilddateien ab, die als notwendige Grundlage für das Einfügen der dreidimensionalen Objektmodelle im folgenden Workflow dienen.
TFRecords und die Spezifikationsdatei
Das TAO Toolkit ist darauf ausgelegt, Daten im TFRecords Format zu verarbeiten, welches das optimierte Binärformat von Googles TensorFlow Framework darstellt. Dieses Format ist notwendig, da es die Trainingsdaten (die generierten Bilder und deren perfekte Annotationen) so effizient speichert und liest, dass die GPU-Auslastung maximiert wird. TFRecords ermöglichen eine schnellere Datenverarbeitung während des Trainings im Vergleich zu einzelnen Bilddateien und CSV-Annotationsdateien, was für das GPU-beschleunigte Transfer Learning auf der RTX 4070 Ti entscheidend ist. Um die Konvertierung der Bilder in dieses spezialisierte Binärformat durchzuführen, wird eine Spezifikationsdatei (Spec-Datei), in diesem Fall convert_spec.txt, benötigt. Spec ist die Abkürzung für Spezifikation und bezeichnet eine einfache Textkonfigurationsdatei, die dem TAO Launcher und dem Konvertierungsskript genau mitteilt, welche Dateien wohin kopiert werden sollen, welches das Quellformat ist (z.B. Bilder und Annotationen) und welche Parameter für die Umwandlung in TFRecords gelten. Die Spec-Datei dient somit als zentrale Steuerungsanweisung für den Konvertierungsprozess, um eine fehlerfreie und reproduzierbare Umwandlung der Quelldaten in das optimierte TFRecords Format zu gewährleisten.
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/
Der Befehl führt die Konvertierung der vorbereiteten Daten in das TFRecords Format aus. Das Ergebnis dieser Konvertierung ist die Erzeugung mehrerer optimierter Binärdateien im TFRecords Format. Diese Dateien werden im angegebenen Ausgabeverzeichnis auf dem Host System gespeichert, welches im Container als /workspace/tao experiments/data/tfrecords/kitti train und auf der lokalen Festplatte als das entsprechend gemountete Verzeichnis vorliegt. Jede dieser Binärdateien enthält die gesamten Trainingsinformationen in einem für TensorFlow optimierten Stream, einschließlich der Bilddaten die synthetischen grauen Bilder mit den gerenderten Objekten sowie deren Annotationen die perfekt annotierten Bounding Boxes und Labels. Das Ergebnis dieses Schrittes schließt die Phase der Datenvorbereitung ab, da die erzeugten TFRecords Dateien die endgültige Trainingsgrundlage für das YOLOv4 Modell sind. Dieses optimierte Format kann in den nächsten Schritten des TAO Workflows vom Trainingsalgorithmus extrem schnell und effizient von der Festplatte gelesen werden, was eine maximale GPU Auslastung gewährleistet.
Training (YOLOv4)
YOLOv4 steht für You Only Look Once Version 4 und repräsentiert eine hochentwickelte Architektur für die Echtzeit Objekterkennung. Im Gegensatz zu älteren Verfahren führt YOLOv4 die Klassifizierung des Objekts und die Lokalisierung der Bounding Box in einem einzigen Durchlauf durch das neuronale Netz aus. Dies ermöglicht die extrem hohe Verarbeitungsgeschwindigkeit die für industrielle Anwendungen wie die fehlerhafte Objekterkennung am Fließband notwendig ist. Für diesen Anwendungsfall nutzt das TAO Toolkit das Modell mit dem ResNet18 als Backbone Architektur, eine kleinere und effizientere Netzwerkstruktur die sich ideal für das Transfer Learning eignet, da es bereits gut auf allgemeinen Bilddaten vor trainiert wurde. Die erstellte Datei yolo_v4_train_resnet18_kitti.txt ist die Hauptkonfigurationsdatei oder Spec Datei für den TAO Launcher und ist für das Training des YOLOv4 Modells absolut unverzichtbar. Sie definiert alle Hyperparameter und Pfade die der Container während des Trainingsprozesses benötigt. Die Konfiguration legt fest welche Lernrate verwendet wird, wie groß die Batch Size ist, wie viele Epochen trainiert werden soll und welche Optimierungsmethode zum Einsatz kommt. Darüber hinaus weist die Spec Datei dem TAO Toolkit an wo die TFRecords Daten liegen die im vorherigen Schritt erzeugt wurden und in welchem Ordner die Trainingsergebnisse gespeichert werden sollen. Sie steuert somit den gesamten Transfer Learning Prozess von der Dateneingabe bis zur Speicherung des trainierten Modells.
Spec-Datei erstellen
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"
}
}
Dateien umbenennen (Fix für Validierung)
Der Konverter erstellt Dateien namens …fold-000…, die Spec-Datei sucht aber nach kitti_val. Wir benennen sie um:
cd ~/tao_workspace/data/tfrecords
for f in *fold-000*; do mv "$f" "${f/kitti_train/kitti_val}"; done
cd ~/tao_workspace
Start des Transfer Learnings
Nachdem alle Vorbereitungen getroffen sind wird nun der Befehl ausgeführt der das eigentliche Transfer Learning auf der NVIDIA GPU startet. Dieser Befehl ist der Höhepunkt des bisherigen Workflows und weist den TAO Launcher an das zuvor heruntergeladene und optimierte YOLOv4 Docker Image zu starten. Dabei wird die erstellte Spec Datei yolo_v4_train_resnet18_kitti.txt als zentrale Konfiguration übergeben. Der Launcher liest diese Spezifikation und mountet die TFRecords Trainingsdaten und die Ergebnisse Ordner in den Container. Anschließend wird innerhalb des Containers das Training initialisiert. Das vortrainierte ResNet18 Backbone des YOLOv4 Modells wird geladen und beginnt die Gewichte basierend auf den synthetisch erzeugten und TFRecords optimierten Daten anzupassen. Die gesamte Rechenlast wird dabei von der RTX 4070 Ti übernommen, wobei das NVIDIA Container Toolkit die effiziente Kommunikation zwischen dem Deep Learning Framework und der GPU sicherstellt. Während des Trainings werden in regelmäßigen Abständen Checkpoint Dateien erzeugt und im gemounteten Ergebnisse Ordner gespeichert, welche den aktuellen Lernfortschritt des Modells darstellen. Diese Checkpoints dienen später zur Evaluierung und zum Export des finalen Modells. Mit diesem Befehl beginnt die Phase in der das Modell lernt defekte von korrekten Objekten zuverlässig zu unterscheiden.
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
Erwarteter Output:
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
Export des Modells in das ONNX Format
Nach dem erfolgreichen Abschluss des Transfer Learnings muss das trainierte YOLOv4 Modell in ein Format überführt werden das für den Produktionseinsatz geeignet ist. Der Befehl Export dient genau diesem Zweck. Er weist den TAO Launcher an das erzeugte Modell zu laden und in das ONNX Format zu konvertieren. ONNX Open Neural Network Exchange ist ein offener Standard, der den Austausch von Deep Learning Modellen zwischen verschiedenen Frameworks und Laufzeitumgebungen ermöglicht. Durch den Export in ONNX wird das Modell frameworkunabhängig gemacht und kann in Engines wie beispielsweise NVIDIA TensorRT geladen werden. TensorRT optimiert das Modell zusätzlich für maximale Leistung und geringste Latenz auf NVIDIA GPUs. Dies ist ein entscheidender Schritt für die Nutzung in Echtzeit Anwendungen wie NVIDIA DeepStream das häufig für Videoanalyse und Bildverarbeitung in der Produktion eingesetzt wird. Das Ergebnis des Exportbefehls ist eine finale Modelldatei im ONNX Format die für die Integration in das Zielsystem bereit ist und die Phase der Modellbereitstellung einleitet.
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
Nachdem die gesamte lokale KI Trainingspipeline erfolgreich eingerichtet wurde, kann der Anwender eigene Projekte starten. Hierzu ist lediglich ein Austausch der synthetischen Bilder gegen eigene Fotos erforderlich. Für die Kennzeichnung der Objekte auf diesen eigenen Fotos wird anschließend die Nutzung eines externen Tools wie beispielsweise labelImg für die Annotationen empfohlen.