ONNX și TensorRT: Optimizarea Modelelor AI pentru Implementarea la Marginea Tactică

Un model care atinge acuratețe de ultimă generație pe un set de date de referință nu se califică automat pentru implementarea pe teren. Modelele de cercetare sunt antrenate și evaluate pe clustere GPU cu memorie și calcul abundente. Modelele implementate trebuie să ruleze pe hardware edge cu bugete de putere stricte, memorie limitată și cerințe de latență în timp real. Decalajul dintre un checkpoint de antrenament PyTorch și un motor de inferență edge de producție acoperă mai mulți pași de optimizare — și fiecare pas are implicații pentru acuratețe, latență și mentenabilitate.

ONNX (Open Neural Network Exchange) și TensorRT sunt cele două tehnologii cheie din pipeline-ul de optimizare a modelelor pentru implementarea NVIDIA Jetson. ONNX furnizează o reprezentare intermediară neutră față de framework care elimină dependența dintre framework-ul de antrenament și mediul de implementare. TensorRT compilează și optimizează modelele ONNX în motoare de inferență specifice dispozitivului care extrag performanța maximă din hardware-ul GPU NVIDIA. Împreună, formează un pipeline de optimizare reproductibil și versionat de la modelul de cercetare la inferența implementată pe teren.

Decalajul antrenament-inferență

Modelul tipic de cercetare există ca un state dict PyTorch (fișier `.pt`) sau TensorFlow SavedModel — o colecție de ponderi învățate și o definiție a grafului de calcul care depinde de runtime-ul framework-ului de antrenament pentru execuție. Runtime-ul framework-ului de antrenament este proiectat pentru flexibilitate (grafuri de calcul dinamice, autograd, gradient checkpointing) mai degrabă decât performanță. Rularea unui model de cercetare în PyTorch pe un NVIDIA Jetson Orin NX fără optimizare produce inferență la aproximativ 8–15 fps pentru un model de detecție de complexitate medie — dramatic sub obiectivul de 30 fps pentru analiza video în timp real.

Trei surse ale acestui decalaj: Overhead de memorie. Checkpoint-urile de antrenament includ starea optimizatorului, buffer-ele pentru gradient și statistici de normalizare batch per strat care nu sunt necesare pentru inferență. Un checkpoint de antrenament YOLOv8-medium este de aproximativ 850 MB; ponderile modelului numai pentru inferență sunt de aproximativ 50 MB. Overhead de calcul. Grafurile de execuție dinamice evaluate la rulare nu pot fi pre-compilate în planuri de execuție statice optimizate pentru hardware. Nepotrivire de precizie. Antrenamentul folosește FP32 sau FP16 pentru precizie numerică; hardware-ul edge atinge randamentul maxim cu INT8, dar conversia în INT8 necesită calibrare pentru a determina scalele de cuantizare adecvate.

ONNX ca schimb universal: export din PyTorch

ONNX definește o reprezentare standard bazată pe graf a arhitecturilor rețelelor neurale, independentă de orice framework de antrenament particular. Exportul unui model PyTorch antrenat în ONNX produce un fișier `.onnx` care conține graful complet de calcul pentru inferență — toate straturile, conexiunile lor și ponderile învățate — într-un format pe care orice runtime compatibil ONNX îl poate executa.

Funcția de export ONNX PyTorch (`torch.onnx.export`) urmărește graful de calcul al modelului rulând cu tensori de intrare exemplu și înregistrând operațiunile rezultate. Principalii parametri de export sunt:

opset_version: Versiunea setului de operatori ONNX. Versiunile mai mari suportă mai multe tipuri de operatori; TensorRT 8.x suportă până la opset 17. Utilizați opset 16 sau 17 pentru acoperire maximă a operatorilor cu versiunile actuale TensorRT.

dynamic_axes: Specifică ce dimensiuni ale tensorului ar trebui tratate ca lungime variabilă (dinamice) mai degrabă decât fixe. Pentru un model de detecție care procesează cadre video, dimensiunea batch ar trebui să fie dinamică pentru a suporta atât inferența în timp real pe un singur cadru (batch=1), cât și procesarea în batch a mai multor cadre. Setarea dimensiunii batch dinamice crește generalitatea modelului ONNX, dar poate împiedica unele optimizări TensorRT care necesită forme fixe.

Obstacole comune la export. Nu toate operațiunile PyTorch au echivalente ONNX directe. Funcțiile autograd personalizate, fluxul de control Python în metodele forward ale modelului (ramuri if/else care depind de valorile tensorilor) și operațiunile adăugate în versiunile recente PyTorch dar care nu sunt încă în setul de operatori ONNX pot provoca eșecuri la export sau grafuri exportate incorecte. Suprimarea non-maximelor din capul de detecție YOLOv8 este cea mai frecventă zonă problematică — Ultralytics furnizează indicatori de configurare a exportului pentru a utiliza implementări NMS compatibile ONNX. Validați întotdeauna modelul ONNX exportat față de modelul PyTorch original pe un set de testare reprezentativ înainte de a continua cu compilarea TensorRT.

Compilarea TensorRT: calibrare INT8, fuziunea straturilor, auto-ajustarea kernelului

TensorRT ia un model ONNX ca intrare și produce un fișier motor TensorRT (`.trt` sau `.engine`) optimizat pentru o arhitectură GPU țintă specifică. Procesul de compilare are trei etape:

Parsarea și validarea rețelei. Parserul ONNX al TensorRT citește graful ONNX și validează că toți operatorii sunt suportați. Operatorii nesuportați trebuie implementați ca plugin-uri personalizate TensorRT înainte de a putea continua compilarea. Pentru arhitecturile CNN standard (variante YOLO, ResNet, MobileNet, EfficientDet), toți operatorii sunt suportați nativ în TensorRT 8.x.

Treceri de optimizare. TensorRT aplică o secvență de optimizări la nivel de graf: fuziunea straturilor (combinând conv+BN+ReLU într-o singură operație), optimizarea aspectului tensorilor (alegând aspectele de memorie NCHW față de NHWC pe baza vitezei pentru fiecare strat pe GPU-ul țintă) și substituția operatorilor (înlocuind anumite secvențe de operații cu alternative echivalente dar mai rapide disponibile ca kernel-uri CUDA pre-construite).

Calibrarea preciziei și compilarea motorului. Pentru compilarea INT8, TensorRT rulează o procedură de calibrare: execută modelul pe un set de date de calibrare reprezentativ (de obicei 500–1.000 de imagini), măsoară intervalul dinamic al activărilor la fiecare strat și determină factorii optimi de scalare a cuantizării. Calitatea datelor de calibrare determină direct acuratețea INT8 — utilizați imagini potrivite domeniului care se potrivesc tipului de senzor de implementare și claselor țintă.

După calibrare, TensorRT evaluează mai multe implementări ale kernel-ului CUDA pentru fiecare strat pe dispozitivul țintă și selectează cea mai rapidă. Acest pas de auto-ajustare poate dura 30–90 de minute pentru un model mare, dar rulează o singură dată — fișierul motor rezultat este serializat și poate fi implementat direct pe dispozitivul țintă fără re-compilare. Fișierele motorului sunt specifice arhitecturii dispozitivului: un motor compilat pe un SKU Jetson (de ex., Orin NX) nu va rula corect pe un SKU diferit (de ex., AGX Orin) din cauza numărătorilor diferite de nuclee GPU.

Latență vs. randament: dimensiunea batch 1 pentru inferența în timp real

Obiectivul de optimizare al TensorRT poate fi configurat pentru latență minimă sau randament maxim. Pentru analiza video în timp real — unde fiecare cadru trebuie procesat și rezultatele livrate înainte de sosirea cadrului următor — latența minimă la dimensiunea batch 1 este obiectivul de optimizare corect. Pentru procesarea offline în batch a imaginilor stocate, randamentul maxim la o dimensiune batch mai mare (8–32) reduce timpul de procesare per cadru cu prețul creșterii latenței per batch.

La dimensiunea batch 1, YOLOv8-medium compilat în INT8 pe Jetson AGX Orin rulează la aproximativ 1,8ms latență (555 fps maxim teoretic). La dimensiunea batch 8, latența crește la aproximativ 7ms per batch (8 × 0,875ms per cadru), dar atinge o eficiență de utilizare a GPU-ului cu aproximativ 15% mai mare. Pentru un flux de intrare la 30fps, dimensiunea batch 1 cu o latență de procesare de 1,8ms oferă marjă substanțială pentru pipeline-uri cu mai multe modele. Pentru un sistem cu patru camere fiecare la 30fps, dimensiunea batch 4 permite procesarea simultană a unui cadru din fiecare flux de cameră în aproximativ 5ms — o arhitectură practică pentru platformele cu mai mulți senzori.

Versionarea modelelor și pipeline-ul de actualizare pentru dispozitivele de teren

Un sistem AI edge implementat necesită un pipeline mentenabil de actualizare a modelelor. Modelele se îmbunătățesc prin date de antrenament suplimentare, rafinamente arhitecturale sau adaptare la noi medii operaționale (teren de iarnă față de vară, noi clase de ținte). Dispozitivele implementate pe teren trebuie să primească aceste actualizări în mod fiabil fără a necesita recuperarea fizică a hardware-ului.

Pipeline-ul de actualizare pentru modelele implementate cu TensorRT diferă de actualizările software standard într-un aspect critic: compilarea TensorRT trebuie executată pe dispozitivul țintă (sau un dispozitiv identic) deoarece include auto-ajustarea kernel-ului specifică dispozitivului. Un flux de lucru pentru actualizarea modelelor ar putea proceda astfel: model nou antrenat la instalație centrală → exportat în ONNX → model ONNX împins pe dispozitiv prin canal de actualizare securizat → compilarea TensorRT executată pe dispozitiv (în timpul unei ferestre de mentenanță când pipeline-ul AI este oprit) → motor compilat activat → motor vechi arhivat pentru rollback.

Fiecare lansare a modelului ONNX ar trebui să conțină un identificator de versiune semantică, un hash de conținut SHA-256 și o semnătură criptografică de la autoritatea de proveniență a modelului. Handlerul de actualizare de pe dispozitiv validează semnătura înainte de execuția compilării. Jurnalul de compilare și hash-ul motorului rezultat ar trebui raportate înapoi la sistemul de management al flotei pentru a confirma actualizarea reușită pe toate dispozitivele din flotă.