ONNX dan Model Export

Paket produksi pada chapter sebelumnya mengikat pemuat pada satu hal: kode Python. state_dict hanya bermakna bagi proses yang mengimpor PyTorch dan definisi class modelnya — sebuah layanan C++, plugin game, atau aplikasi ponsel tidak bisa membacanya. ONNX (Open Neural Network Exchange) menjawab batas ini dengan cara memindahkan model keluar dari framework sama sekali: arsitektur dan bobot dibekukan menjadi graph computation standar, lalu dijalankan oleh runtime apa pun yang memahami format itu. Chapter ini mengekspor model PyTorch ke ONNX, menjalankannya dengan onnxruntime, dan memverifikasi kesetaraan outputnya secara nyata di mesin ini.

Graph Standar dan Runtime Terpisah #

File .pth menyimpan parameter; file .onnx menyimpan program. Sebuah model ONNX adalah graph berarah yang node-nya adalah operasi matematika baku (Gemm, Relu, Add) dan tepinya adalah tensor bernama, lengkap dengan bobot yang sudah tertanam sebagai constant. Runtime seperti onnxruntime membaca graph itu, mengoptimasinya, lalu mengeksekusinya — tanpa PyTorch, tanpa Python, tanpa kode class kita.

Sejak PyTorch 2.9, torch.onnx.export memakai exporter baru berbasis torch.export (jalur dynamo) sebagai bawaan — selaras dengan status TorchScript yang telah deprecated di 2.14. Parameter penting yang berganti wajah: dimensi dinamis kini dinyatakan lewat dynamic_shapes, bukan lagi dynamic_axes gaya lama.

Praktik: Export MLP dan Inference dengan onnxruntime #

Berkas export_onnx.py melatih MLP pada dataset make_moons, mengekspornya dengan dimensi batch dinamis, lalu membandingkan output PyTorch dan onnxruntime pada seluruh data uji:

"""Export model PyTorch ke ONNX dan inference dengan onnxruntime.

Alur: latih MLP kecil -> export ke model.onnx lewat torch.onnx.export
(exporter dynamo bawaan PyTorch 2.14) -> muat dengan onnxruntime ->
bandingkan output torch vs onnxruntime pada seluruh data uji ->
laporkan selisih maksimum -> bersihkan berkas demo.
"""

from pathlib import Path

import numpy as np
import torch
import torch.nn as nn
from sklearn.datasets import make_moons
from sklearn.model_selection import train_test_split
from torch.utils.data import DataLoader, TensorDataset

torch.manual_seed(42)
X, y = make_moons(n_samples=256, noise=0.20, random_state=42)
X = (X - X.mean(axis=0)) / X.std(axis=0)
X_latih, X_uji = train_test_split(X, test_size=0.25, random_state=42)

loader = DataLoader(
    TensorDataset(
        torch.tensor(X_latih, dtype=torch.float32),
        torch.tensor(train_test_split(y, test_size=0.25, random_state=42)[0],
                     dtype=torch.long),
    ),
    batch_size=32,
    shuffle=True,
    num_workers=0,
    generator=torch.Generator().manual_seed(42),
)
model = nn.Sequential(nn.Linear(2, 16), nn.ReLU(), nn.Linear(16, 2))
kriteria = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.5)
for _ in range(20):
    for X_batch, y_batch in loader:
        optimizer.zero_grad()
        loss = kriteria(model(X_batch), y_batch)
        loss.backward()
        optimizer.step()
model.eval()

# ---------- 1. Export ke ONNX ----------
contoh = torch.tensor(X_uji[:1], dtype=torch.float32)
torch.onnx.export(
    model,
    (contoh,),
    "model.onnx",
    input_names=["fitur"],
    output_names=["logits"],
    dynamic_shapes={"input": {0: "jumlah_sampel"}},
    verbose=False,
)
print("export selesai -> model.onnx (dimensi batch dinamis)")

# ---------- 2. Inference dengan onnxruntime ----------
import onnxruntime as ort

sesi = ort.InferenceSession("model.onnx", providers=["CPUExecutionProvider"])
print("provider onnxruntime :", sesi.get_providers())

X_uji_tensor = torch.tensor(X_uji, dtype=torch.float32)
with torch.no_grad():
    logits_torch = model(X_uji_tensor).numpy()
logits_ort = sesi.run(None, {"fitur": X_uji_tensor.numpy()})[0]

selisih = np.abs(logits_torch - logits_ort)
cocok = np.allclose(logits_torch, logits_ort, atol=1e-5)
print(f"jumlah sampel uji    : {len(X_uji)}")
print(f"selisih maksimum     : {selisih.max():.2e}")
print(f"output identik <1e-5 : {cocok}")
print(f"prediksi kelas sama  : "
      f"{(logits_torch.argmax(1) == logits_ort.argmax(1)).mean():.4f}")

ukuran_torch = Path("model.onnx").stat().st_size
print(f"ukuran berkas onnx   : {ukuran_torch} byte")

Path("model.onnx").unlink()
print("berkas model.onnx dihapus dari bundle")

Output yang Dihasilkan:

export selesai -> model.onnx (dimensi batch dinamis)
provider onnxruntime : ['CPUExecutionProvider']
jumlah sampel uji    : 64
selisih maksimum     : 0.00e+00
output identik <1e-5 : True
prediksi kelas sama  : 1.0000
ukuran berkas onnx   : 4601 byte
berkas model.onnx dihapus dari bundle

Mari kita telaah kode di atas baris per baris:

1. Baris 46-55 (torch.onnx.export): contoh input satu sampel cukup sebagai bentuk acuan — graph direkam dari eksekusi nyata model. input_names dan output_names memberi nama tensor di graph sehingga pemanggil runtime bisa menyebut "fitur" alih-alih indeks anonim.
2. Baris 53 (dynamic_shapes): kunci "input" adalah nama argumen forward model (nama bawaan nn.Sequential), bukan nama ONNX; nilai {0: "jumlah_sampel"} menyatakan dimensi ke-0 bebas nilainya. Berkat ini, satu berkas model.onnx melayani batch 1 maupun batch 64 — dibuktikan keluaran yang menguji 64 sampel sekaligus.
3. Baris 57-60 (InferenceSession): onnxruntime memuat graph sekali dan mengoptimasinya; providers=["CPUExecutionProvider"] secara eksplisit memilih jalur CPU. Di mesin dengan GPU, daftar provider bisa ditambah CUDAExecutionProvider — urutannya menentukan prioritas.
4. Baris 68-69 (sesi.run): pemanggilan inference menerima NumPy array dan nama input, mengembalikan daftar output. Tidak ada PyTorch di jalur ini sama sekali — itulah inti portabilitas ONNX.
5. Baris 71-79 (verifikasi): selisih maksimum 0.00e+00 berarti kedua jalur menghasilkan bit yang sama persis pada 64 sampel uji — hasil terbaik yang bisa diharapkan dari export; pada model yang lebih kompleks, selisih kecil orde 1e-6 tetap wajar karena perbedaan urutan operasi floating point.

Dokumentasi lengkap runtime, provider, dan opsi performanya tersedia di onnxruntime.ai/docs , sedangkan API torch.onnx di PyTorch 2.14 dibahas di docs.pytorch.org/docs/stable/onnx.html , dan spesifikasi formatnya di onnx.ai .

Model yang sudah menjadi graph terbuka peluang optimasi baru: graph itu bisa disusun ulang, digabung operasinya, dan dijalankan dengan pengaturan thread yang berbeda dari framework asalnya. Chapter berikutnya mengambil sumber latensi satu per satu dan mempraktikkan dua pengungkitnya — pengaturan thread CPU dan graph optimization onnxruntime — di mesin dua core ini.