Bài 56: ONNX và TensorRT — export model cho inference tối ưu

Sau bài này bạn sẽ:

• Hiểu ONNX là format gì và tại sao cần nó trong deployment pipeline
• Export model PyTorch và Hugging Face sang ONNX đúng cách với dynamic_axes
• Chạy ONNX model bằng ONNX Runtime trên CPU và CUDA
• Compile ONNX sang TensorRT engine FP16/INT8
• Biết khi nào dùng ONNX Runtime, TensorRT, OpenVINO, CoreML
• Nắm các pitfalls hay gặp khi export và deploy

Open Neural Network Exchange (ONNX) là format file chuẩn mở để lưu model neural network. Microsoft, Meta và Amazon khởi xướng năm 2017; hiện có hơn 50 framework và runtime hỗ trợ.

ONNX mô tả model dưới dạng computation graph: mỗi node là một operator (Conv, MatMul, ReLU…), mỗi edge là tensor. File .onnx là protobuf binary chứa graph + weight.

Framework có thể export sang ONNX

• PyTorch — torch.onnx.export() (built-in)
• TensorFlow / Keras — tf2onnx
• Hugging Face Transformers — optimum CLI và Python API
• scikit-learn — sklearn-onnx
• XGBoost, LightGBM — onnxmltools

Runtime có thể đọc ONNX

• ONNX Runtime — cross-platform, CPU và GPU (CUDA, DirectML, ROCm)
• TensorRT — NVIDIA GPU
• OpenVINO — Intel CPU, iGPU, VPU
• CoreML — Apple Silicon, iOS, macOS
• NNAPI / TFLite — Android

Khi nào cần ONNX

Trường hợp phổ biến nhất: train với PyTorch, nhưng muốn serve bằng runtime tối ưu hơn (ONNX Runtime, TensorRT) hoặc deploy trên hardware không có PyTorch (edge device, mobile, cloud inference endpoint tối thiểu hóa dependency).

TensorRT là SDK của NVIDIA để tối ưu và chạy deep learning inference trên GPU NVIDIA. TensorRT không phải runtime đơn thuần — nó là compiler: nhận ONNX (hoặc model PyTorch, TF) rồi tạo ra engine binary tối ưu cho GPU cụ thể.

Cơ chế TensorRT hoạt động

• Layer fusion: gộp nhiều operator liền nhau (Conv + BN + ReLU) thành 1 kernel, giảm roundtrip bộ nhớ GPU
• Precision calibration: hỗ trợ FP32, FP16, INT8 — INT8 cần calibration data
• Kernel auto-tuning: chọn kernel implementation tốt nhất cho GPU cụ thể (T4, A100, H100…)
• Memory optimization: tái sử dụng tensor buffer trong graph

Trade-off cần biết trước

• Engine binary lock vào GPU type: compile trên T4 thì chỉ chạy trên T4; deploy GPU khác phải compile lại
• Build engine mất vài phút đến hàng chục phút tùy model và GPU
• Chỉ chạy trên NVIDIA GPU — không có fallback CPU
• Hỗ trợ dynamic shape nhưng phải khai báo tường minh trong config

TensorRT 10.x (phát hành 2024) có Python API ổn định hơn; các ví dụ dưới dùng TensorRT 10.x.

PyTorch model
    │
    ├─── torch.onnx.export() ──→ model.onnx
    │                                │
    │                                ├──→ ONNX Runtime (CPU / CUDA / DirectML)
    │                                │
    │                                ├──→ TensorRT engine (NVIDIA GPU only)
    │                                │
    │                                ├──→ OpenVINO (Intel CPU / iGPU)
    │                                │
    │                                └──→ CoreML (Apple Silicon, iOS)
    │
    └─── TensorRT Python API ──→ TensorRT engine (direct, không qua ONNX)

Con đường phổ biến nhất trong production là PyTorch → ONNX → ONNX Runtime hoặc TensorRT. Bước ONNX trung gian giúp decouple training framework khỏi serving runtime.

import torch
import torchvision.models as models

model = models.resnet18(weights="IMAGENET1K_V1").eval()
dummy_input = torch.randn(1, 3, 224, 224)

torch.onnx.export(
    model,
    dummy_input,
    "resnet18.onnx",
    export_params=True,       # lưu weight vào file
    opset_version=17,         # ONNX opset — dùng 17+ cho model 2023 trở về sau
    do_constant_folding=True, # fold constant operation lúc export
    input_names=["input"],
    output_names=["output"],
    dynamic_axes={
        "input":  {0: "batch_size"},  # batch dim là variable
        "output": {0: "batch_size"},
    },
)
print("Exported resnet18.onnx")

Các tham số quan trọng

• opset_version: ONNX định nghĩa operator theo opset. Opset 17 là stable cho hầu hết model 2024. Nếu runtime báo lỗi operator không hỗ trợ, thử hạ xuống 13 hoặc 15.
• dynamic_axes: không khai báo thì engine sẽ fix shape = shape của dummy_input. Trong production gần như luôn cần batch variable — khai báo tường minh.
• do_constant_folding: pre-compute node có output là constant (vd bias, BN statistics) lúc export. Giảm nhẹ graph size.

Verify file vừa export

import onnx

model_onnx = onnx.load("resnet18.onnx")
onnx.checker.check_model(model_onnx)
print(f"ONNX opset: {model_onnx.opset_import[0].version}")
print(f"Input:  {[i.name for i in model_onnx.graph.input]}")
print(f"Output: {[o.name for o in model_onnx.graph.output]}")

onnx.checker.check_model() raise ValidationError nếu graph không hợp lệ. Chạy bước này trước khi tiếp tục.

Transformer model có dynamic input (token length variable) và nhiều output — export thủ công với torch.onnx.export khá phức tạp. Thư viện Optimum của Hugging Face tự động hóa bước này.

Cài đặt

pip install "optimum[exporters]"   # optimum 1.20+

Export qua CLI

# Export DistilBERT cho classification
optimum-cli export onnx \
  --model distilbert-base-uncased-finetuned-sst-2-english \
  --task text-classification \
  distilbert_onnx/

# Export BERT cho feature extraction
optimum-cli export onnx \
  --model bert-base-uncased \
  --task feature-extraction \
  bert_onnx/

Flag --task quan trọng — Optimum dùng nó để biết cần bao gồm head nào trong graph.

Export và chạy ngay trong Python

from optimum.onnxruntime import ORTModelForSequenceClassification
from transformers import AutoTokenizer

# Export khi load lần đầu (export=True), lưu vào thư mục
model = ORTModelForSequenceClassification.from_pretrained(
    "distilbert-base-uncased-finetuned-sst-2-english",
    export=True,
)
model.save_pretrained("./distilbert_onnx")

tokenizer = AutoTokenizer.from_pretrained(
    "distilbert-base-uncased-finetuned-sst-2-english"
)
inputs = tokenizer("I love this movie", return_tensors="pt")
outputs = model(**inputs)
print(outputs.logits)

Lần sau chỉ cần from_pretrained("./distilbert_onnx") — không cần export=True nữa. ORTModelForXxx có interface giống AutoModelForXxx, drop-in replacement trong pipeline.

import onnxruntime as ort  # onnxruntime 1.18+
import numpy as np

# Khai báo provider theo thứ tự ưu tiên
# ONNX Runtime sẽ dùng provider đầu tiên available
session = ort.InferenceSession(
    "resnet18.onnx",
    providers=["CUDAExecutionProvider", "CPUExecutionProvider"],
)

# Kiểm tra provider nào đang được dùng
print(session.get_providers())

# Lấy tên input/output từ session
input_name  = session.get_inputs()[0].name   # "input"
output_name = session.get_outputs()[0].name  # "output"

input_data = np.random.randn(1, 3, 224, 224).astype(np.float32)
outputs = session.run([output_name], {input_name: input_data})
print(outputs[0].shape)  # (1, 1000)

Execution Providers

• CUDAExecutionProvider — NVIDIA GPU, cần onnxruntime-gpu
• TensorrtExecutionProvider — NVIDIA GPU qua TensorRT, cần cả TensorRT cài sẵn
• CPUExecutionProvider — CPU, luôn có, là fallback
• OpenVINOExecutionProvider — Intel hardware
• CoreMLExecutionProvider — Apple Silicon, macOS / iOS

Nếu CUDAExecutionProvider không available (không có GPU hoặc cài onnxruntime thay vì onnxruntime-gpu), ONNX Runtime tự fallback sang CPUExecutionProvider. Không báo lỗi — cần log session.get_providers() để confirm.

Session Options để tune

opts = ort.SessionOptions()
opts.intra_op_num_threads = 4  # số thread cho 1 op
opts.inter_op_num_threads = 1  # số thread song song giữa các op
opts.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL

session = ort.InferenceSession(
    "resnet18.onnx",
    sess_options=opts,
    providers=["CUDAExecutionProvider", "CPUExecutionProvider"],
)

Sau khi export, output của ONNX Runtime phải khớp với PyTorch trong tolerance số học (FP32 có sai số nhỏ do thứ tự operation).

import torch
import numpy as np

# PyTorch forward pass
dummy_input = torch.randn(1, 3, 224, 224)
with torch.no_grad():
    torch_output = model(dummy_input).numpy()

# ONNX Runtime forward pass
ort_output = session.run(
    [output_name],
    {input_name: dummy_input.numpy()},
)[0]

# So sánh
np.testing.assert_allclose(
    torch_output,
    ort_output,
    rtol=1e-3,  # relative tolerance
    atol=1e-5,  # absolute tolerance
)
print("Output matches within tolerance")

# Xem max diff để đánh giá
max_diff = np.abs(torch_output - ort_output).max()
print(f"Max absolute difference: {max_diff:.2e}")

Nếu diff lớn hơn expected (>1e-2): thường do mixed precision trong model hoặc custom operation. Thử tắt do_constant_folding hoặc đổi opset_version.

Với INT8 quantized ONNX model, tolerance cần nới rộng hơn (rtol=1e-1) — accuracy giảm là expected. Cần eval trên validation set thay vì chỉ so sánh output 1 sample.

import time
import torch
import numpy as np

def benchmark(forward_fn, input_data, n: int = 100, warmup: int = 10) -> float:
    """Trả về latency trung bình (ms) trên n lần chạy."""
    for _ in range(warmup):
        forward_fn(input_data)

    start = time.perf_counter()
    for _ in range(n):
        forward_fn(input_data)
    elapsed = (time.perf_counter() - start) / n * 1000
    return elapsed

input_np = np.random.randn(1, 3, 224, 224).astype(np.float32)

# PyTorch CPU
with torch.no_grad():
    torch_latency = benchmark(
        lambda x: model(torch.from_numpy(x)),
        input_np,
    )

# ONNX Runtime
ort_latency = benchmark(
    lambda x: session.run(None, {input_name: x}),
    input_np,
)

print(f"PyTorch CPU:   {torch_latency:.2f} ms")
print(f"ONNX Runtime:  {ort_latency:.2f} ms")
print(f"Speedup:       {torch_latency / ort_latency:.2f}x")

Lưu ý khi benchmark GPU: cần thêm torch.cuda.synchronize() (PyTorch) và xài CUDA event để đo chính xác, không phải time.perf_counter(). CPU benchmark với time.perf_counter() là đủ.

Trên CPU, ONNX Runtime thường nhanh hơn PyTorch 1.5–2x với model classification. Khoảng cách lớn hơn nếu model nhiều element-wise operation (ONNX Runtime fuse tốt hơn).

Cách nhanh nhất: trtexec CLI

# FP32
trtexec --onnx=resnet18.onnx --saveEngine=resnet18_fp32.engine

# FP16 — gần như không mất accuracy, speedup ~1.5-2x so với FP32
trtexec --onnx=resnet18.onnx --saveEngine=resnet18_fp16.engine --fp16

# Dynamic shape — khai báo minShapes/optShapes/maxShapes
trtexec \
  --onnx=resnet18.onnx \
  --saveEngine=resnet18_dynamic.engine \
  --fp16 \
  --minShapes=input:1x3x224x224 \
  --optShapes=input:8x3x224x224 \
  --maxShapes=input:32x3x224x224

trtexec đi kèm với TensorRT installation (/usr/src/tensorrt/bin/trtexec hoặc trong container nvcr.io/nvidia/tensorrt). Đây là cách đủ dùng cho phần lớn trường hợp không cần custom config.

Python API — khi cần kiểm soát chi tiết

import tensorrt as trt  # tensorrt 10.x

TRT_LOGGER = trt.Logger(trt.Logger.WARNING)

builder = trt.Builder(TRT_LOGGER)
network = builder.create_network(
    1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)
)
parser = trt.OnnxParser(network, TRT_LOGGER)

with open("resnet18.onnx", "rb") as f:
    if not parser.parse(f.read()):
        for i in range(parser.num_errors):
            print(parser.get_error(i))
        raise RuntimeError("ONNX parse failed")

config = builder.create_builder_config()
config.set_flag(trt.BuilderFlag.FP16)
# 1GB workspace cho optimization
config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 1 << 30)

# Dynamic shape profile
profile = builder.create_optimization_profile()
profile.set_shape(
    "input",
    min=(1,  3, 224, 224),
    opt=(8,  3, 224, 224),
    max=(32, 3, 224, 224),
)
config.add_optimization_profile(profile)

engine_bytes = builder.build_serialized_network(network, config)
with open("resnet18.engine", "wb") as f:
    f.write(engine_bytes)
print("Engine saved")

Build engine có thể mất 5–20 phút tùy model và GPU. Nên chạy một lần lúc CI/CD, không chạy lúc startup service.

FP16 — gần như miễn phí

Chỉ cần thêm config.set_flag(trt.BuilderFlag.FP16). TensorRT tự quyết định layer nào convert sang FP16, layer nào giữ FP32. Accuracy gần như không giảm trên classification model; speedup 1.5–2x so với FP32 trên GPU có Tensor Core (T4, A100, H100).

INT8 — cần calibration data

INT8 yêu cầu calibrator: một tập data đại diện để TensorRT học range của activation và chọn scale factor tối ưu.

import tensorrt as trt
import numpy as np

class Int8EntropyCalibrator(trt.IInt8EntropyCalibrator2):
    """
    Calibrator dùng IInt8EntropyCalibrator2 (entropy-based,
    thường tốt hơn MinMax cho classification).
    """
    def __init__(self, calib_data: np.ndarray, cache_file: str = "calib.cache"):
        super().__init__()
        self._data   = calib_data          # shape: (N, C, H, W), float32
        self._idx    = 0
        self._cache  = cache_file
        # Allocate GPU buffer
        import pycuda.driver as cuda
        import pycuda.autoinit  # noqa: F401
        self._buf = cuda.mem_alloc(calib_data[0].nbytes)

    def get_batch_size(self) -> int:
        return 1

    def get_batch(self, names):
        if self._idx >= len(self._data):
            return None
        import pycuda.driver as cuda
        cuda.memcpy_htod(self._buf, np.ascontiguousarray(self._data[self._idx]))
        self._idx += 1
        return [int(self._buf)]

    def read_calibration_cache(self):
        if self._cache and open(self._cache, "rb").read() if __import__("os").path.exists(self._cache) else None:
            with open(self._cache, "rb") as f:
                return f.read()

    def write_calibration_cache(self, cache):
        with open(self._cache, "wb") as f:
            f.write(cache)


# Sử dụng
calib_data = np.random.randn(100, 3, 224, 224).astype(np.float32)  # dùng real data
calibrator = Int8EntropyCalibrator(calib_data)

config.set_flag(trt.BuilderFlag.INT8)
config.int8_calibrator = calibrator

Calibration data cần representative: nếu dùng data không đại diện cho distribution production, scale factor sai, accuracy giảm đáng kể. Tối thiểu 100–500 sample, lý tưởng 1000.

TensorRT lưu cache calibration (calib.cache). Lần sau build engine với cùng model thì đọc lại cache — không cần chạy calibration lại từ đầu nếu model không thay đổi.

import tensorrt as trt
import pycuda.driver as cuda
import pycuda.autoinit  # noqa: F401
import numpy as np

# Load engine
runtime = trt.Runtime(trt.Logger(trt.Logger.WARNING))
with open("resnet18.engine", "rb") as f:
    engine = runtime.deserialize_cuda_engine(f.read())

context = engine.create_execution_context()

# Allocate I/O buffers (một lần lúc init)
input_shape  = (1, 3, 224, 224)
output_shape = (1, 1000)

h_input  = np.random.randn(*input_shape).astype(np.float32)
h_output = np.empty(output_shape, dtype=np.float32)

d_input  = cuda.mem_alloc(h_input.nbytes)
d_output = cuda.mem_alloc(h_output.nbytes)
stream   = cuda.Stream()

# Inference
cuda.memcpy_htod_async(d_input, h_input, stream)
context.execute_async_v2(
    bindings=[int(d_input), int(d_output)],
    stream_handle=stream.handle,
)
cuda.memcpy_dtoh_async(h_output, d_output, stream)
stream.synchronize()

print(h_output.argmax())

TensorRT Python API verbose hơn ONNX Runtime vì phải tự quản lý GPU memory. Trong production, gần như không ai viết inference loop này trực tiếp — thay vào đó dùng Triton Inference Server (mục tiếp theo) hoặc wrap bằng thư viện như torch_tensorrt.

torch-tensorrt — ít code hơn

import torch
import torch_tensorrt  # torch-tensorrt 2.x

# Compile trực tiếp từ PyTorch model, không qua ONNX
trt_model = torch_tensorrt.compile(
    model,
    inputs=[torch_tensorrt.Input(
        min_shape=[1,  3, 224, 224],
        opt_shape=[8,  3, 224, 224],
        max_shape=[32, 3, 224, 224],
        dtype=torch.float32,
    )],
    enabled_precisions={torch.float16},
)

# Inference như PyTorch thường
with torch.no_grad():
    output = trt_model(torch.randn(1, 3, 224, 224).cuda())

torch_tensorrt giữ interface PyTorch, compile ngầm sang TensorRT. Phù hợp nếu code base đã dùng PyTorch và không muốn đổi sang pycuda.

Triton Inference Server (NVIDIA) là pattern production chuẩn để serve TensorRT engine — không dùng Python inference loop thủ công như mục trước.

Cấu trúc model repository

model_repository/
└── resnet18/
    ├── config.pbtxt          # cấu hình model
    └── 1/
        └── model.plan        # file .engine đổi tên thành model.plan

# config.pbtxt
name: "resnet18"
platform: "tensorrt_plan"
max_batch_size: 32
input [
  {
    name: "input"
    data_type: TYPE_FP32
    dims: [3, 224, 224]
  }
]
output [
  {
    name: "output"
    data_type: TYPE_FP32
    dims: [1000]
  }
]
dynamic_batching {
  preferred_batch_size: [8, 16]
  max_queue_delay_microseconds: 5000
}

Khởi động server

docker run --gpus all --rm \
  -p 8000:8000 -p 8001:8001 \
  -v $(pwd)/model_repository:/models \
  nvcr.io/nvidia/tritonserver:24.05-py3 \
  tritonserver --model-repository=/models

Client gọi HTTP

import tritonclient.http as httpclient
import numpy as np

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

input_data = np.random.randn(1, 3, 224, 224).astype(np.float32)
inputs  = [httpclient.InferInput("input",  input_data.shape, "FP32")]
outputs = [httpclient.InferRequestedOutput("output")]
inputs[0].set_data_from_numpy(input_data)

result = client.infer("resnet18", inputs, outputs=outputs)
print(result.as_numpy("output").argmax())

Triton tự xử lý dynamic batching (gom request đến trong window max_queue_delay_microseconds thành 1 batch), concurrent execution trên nhiều GPU, và health check endpoint. Đây là lý do production không viết inference server tự tay.

TensorRT thông thường không tối ưu được LLM vì LLM có KV cache, autoregressive generation, và shape thay đổi từng token. NVIDIA cung cấp thư viện riêng: TensorRT-LLM.

• Hỗ trợ sẵn Llama 2/3, Mistral, Falcon, GPT-J, Gemma, Phi… (danh sách mở rộng mỗi release)
• Tích hợp continuous batching và PagedAttention — tương tự vLLM
• Hỗ trợ tensor parallelism cho multi-GPU
• Backend native cho Triton Inference Server

# Build LLaMA 3 8B từ HuggingFace weights
python convert_checkpoint.py \
  --model_dir ./Meta-Llama-3-8B \
  --output_dir ./llama3-trtllm \
  --dtype float16

trtllm-build \
  --checkpoint_dir ./llama3-trtllm \
  --output_dir ./llama3-engine \
  --gemm_plugin float16 \
  --max_batch_size 8 \
  --max_input_len 2048 \
  --max_seq_len 4096

TensorRT-LLM vs vLLM: Cả hai đều dùng PagedAttention. TensorRT-LLM thường nhanh hơn vài phần trăm đến 20% tùy model và batch size; nhưng setup phức tạp hơn, lock vào NVIDIA, compile lại khi đổi GPU. vLLM đơn giản hơn, hỗ trợ AMD GPU, AWS Inferentia. Với team nhỏ không có dedicated MLOps, vLLM thường là lựa chọn thực tế hơn.

Số liệu bên dưới từ benchmark ResNet-50, batch=1, trên T4 và A100 80GB. Kết quả cụ thể sẽ khác theo model architecture và batch size.

Speedup từ FP16 có được nhờ Tensor Core trên GPU NVIDIA từ Volta trở về sau (T4, A100, H100). GPU cũ (P100, K80) không có Tensor Core — FP16 speedup rất nhỏ hoặc không đáng kể.

Các trường hợp khác

• OpenVINO: deploy trên Intel CPU (server không có GPU, edge gateway), iGPU Intel. Thường nhanh hơn ONNX Runtime 1.5–3x trên Intel hardware.
• CoreML: iOS app, macOS app, Apple Silicon. Apple Neural Engine cho speedup đáng kể trên M1/M2/M3.
• GGUF + llama.cpp: LLM inference trên CPU (không cần CUDA), đã bàn trong bài 55.

• Export ONNX với batch cố định. Nếu không khai báo dynamic_axes, engine fix shape theo dummy_input. Khi production cần batch khác phải re-export. Luôn khai báo dynamic_axes cho dimension cần flexible.
• Opset version không tương thích. Opset 11 hoặc 12 đủ cho nhiều model cũ, nhưng TensorRT 10.x và ONNX Runtime 1.18+ khuyến nghị opset 17. Nếu runtime báo lỗi no implementation for operator X at opset Y, kiểm tra lại opset.
• TensorRT engine không portable. Engine compile trên T4 không load được trên A100 và ngược lại. Cần build pipeline riêng cho từng GPU type trong fleet.
• Dynamic shape trong TensorRT chưa khai báo. Nếu ONNX model có dynamic batch nhưng TensorRT config không set optimization_profile, engine sẽ fix shape hoặc báo lỗi lúc inference. Dùng --minShapes/--optShapes/--maxShapes với trtexec hoặc set_shape() trong Python API.
• INT8 calibration data không representative. Scale factor tính từ data không đúng distribution production → activation clip sai → accuracy giảm 5–15%. Dùng ít nhất vài trăm sample từ production data, không dùng random noise.
• Custom PyTorch operator không có trong ONNX opset. Một số custom layer (custom attention variant, custom normalization) không map được sang ONNX standard op. Export sẽ fail hoặc fallback sang ATen op mà runtime không hỗ trợ. Cần register custom op hoặc refactor layer dùng standard op.
• ONNX Runtime trên Mac M1/M2 với CUDAExecutionProvider. Mac không có CUDA; ONNX Runtime fallback CPU nhưng không báo lỗi. Dùng CoreMLExecutionProvider để tận dụng Apple Neural Engine.
• Không verify accuracy sau export. Export thành công không có nghĩa output đúng. Luôn chạy np.testing.assert_allclose với vài sample và eval trên validation set trước khi deploy.
• Build TensorRT engine lúc startup service. Build mất nhiều phút, service sẽ không sẵn sàng trong thời gian đó. Build trong CI/CD pipeline, lưu engine artifact, load pre-built engine khi startup.