Vision Model များကို Deploy လုပ်ခြင်း

နားလည်ထားရမယ့် အချက်

ONNX (Open Neural Network Exchange) ဆိုတာ model တစ်ခုရဲ့ computation graph နဲ့ weight တွေကို train လုပ်ခဲ့တဲ့ framework နဲ့ သီးခြားစီ ဖမ်းယူထားတဲ့ intermediate format တစ်ခုပါ။ PyTorch နဲ့ train လုပ်ထားတဲ့ model တစ်ခုကို ONNX Runtime, TensorRT, mobile inference engine (Android/iOS ပေါ်) ဒါမှမဟုတ် onnxruntime-web ကနေတဆင့် browser ပေါ်မှာတောင် PyTorch install အပြည့်အစုံ မလိုအပ်ဘဲ run နိုင်စေပါတယ်။ torch.onnx.export function က representative dummy input တစ်ခုပေးလိုက်တဲ့အခါ model ရဲ့ forward pass ကို trace (သို့) symbolically capture လုပ်ပြီး operation တစ်ခုချင်းစီကို graph node အနေနဲ့ မှတ်တမ်းတင်ပါတယ်။

လက်တွေ့ scenario နဲ့ ချိတ်ကြည့်မယ်

Tutorial Platform ရဲ့ upload processing pipeline ထဲမှာ diagram ပုံတွေအတွက် caption auto-generate လုပ်ပေးတဲ့ vision model တစ်ခု run နေပါတယ်၊ ဒါပေမယ့် ဒီ model ကို serverless function ထဲမှာ run ဖို့ PyTorch full runtime တစ်ခုလုံးကို bundle လုပ်ရရင် cold-start time ကြာလွန်းနေပါတယ်။ Platform team က ဒီ model ကို ONNX ဆီ export လုပ်ပြီး ONNX Runtime ရဲ့ lightweight package ကိုသာ deploy dependency အဖြစ် အသုံးပြုကာ serverless function ရဲ့ cold-start latency ကို သိသိသာသာ လျှော့ချနိုင်ပါတယ်။

ONNX format က model ကို train လုပ်ခဲ့တဲ့ framework နဲ့ သီးခြားစီထားပြီး inference runtime များစွာ (edge device အပါအဝင်) ကနေ run နိုင်စေပါတယ်။
torch.onnx.export က dummy input တစ်ခုအသုံးပြုကာ model ရဲ့ forward pass ကို graph အနေနဲ့ capture လုပ်ကာ .onnx file အဖြစ် သိမ်းဆည်းပေးပါတယ်။
Quantization (float32 ကနေ int8 ဆီ ပြောင်းခြင်း) က model size ကို လျှော့ချပြီး edge/CPU hardware ပေါ်မှာ inference speed ကို မြှင့်တင်ပေးတတ်ပြီး accuracy ကတော့ ပုံမှန်အားဖြင့် သိသိသာသာ မကျစေပါဘူး။

အတူတူ စမ်းရေးကြည့်မယ်

python

import torch
import torch.nn as nn

torch.manual_seed(0)


class TinyClassifier(nn.Module):
    def __init__(self, num_classes=10):
        super().__init__()
        self.conv = nn.Sequential(
            nn.Conv2d(3, 8, kernel_size=3, stride=2, padding=1),  # 32x32 -> 16x16
            nn.ReLU(),
            nn.AdaptiveAvgPool2d(1),
        )
        self.classifier = nn.Linear(8, num_classes)

    def forward(self, x):
        features = self.conv(x).flatten(1)
        return self.classifier(features)


model = TinyClassifier()
model.eval()  # disable dropout/batchnorm training behavior before export

dummy_input = torch.randn(1, 3, 32, 32)  # one fake RGB image, batch size 1

onnx_path = "tiny_classifier.onnx"
# Requires the onnx and onnxscript packages in addition to torch/torchvision:
#   pip install onnx onnxscript
torch.onnx.export(
    model,
    dummy_input,
    onnx_path,
    input_names=["image"],
    output_names=["class_scores"],
    dynamic_axes={"image": {0: "batch_size"}, "class_scores": {0: "batch_size"}},
    opset_version=18,
    verbose=False,
)

print(f"Exported model to {onnx_path}")
print("This ONNX file can now run on ONNX Runtime, mobile, or edge inference engines")
print("Input shape:", dummy_input.shape)
print("Output shape:", model(dummy_input).shape)

You should see

Export ပြီးဆုံးရင် tiny_classifier.onnx ဆိုတဲ့ file ကို working directory ထဲမှာ ဖန်တီးပါတယ် (torch.onnx.export ရဲ့ internal exporter က export process အတွင်း diagnostic log အချို့ console ပေါ် ထပ်ပြသနိုင်ပါတယ်)။ ပြီးရင် Exported model to tiny_classifier.onnx၊ ONNX file ရဲ့ အသုံးဝင်ပုံအကြောင်း message တစ်ကြောင်း၊ Input shape: torch.Size([1, 3, 32, 32])၊ ပြီးတော့ Output shape: torch.Size([1, 10]) — class 10 ခုအတွက် score vector တစ်ခု — ကို print ထုတ်ပါတယ်။

၅ မိနစ် စမ်းကြည့်

Code ကို ပြင်ပြီး dummy_input ကို batch size 1 အစား torch.randn(4, 3, 32, 32) ဖြင့် ပြောင်းလိုက်ပါ၊ dynamic_axes ကို သတ်မှတ်ထားပြီးသားဖြစ်လို့ export ဆက်လက် အောင်မြင်သင့်ပါတယ်။ Export ပြီးနောက် exported model ရဲ့ output shape ကို batch size 4 အတွက် ဘယ်လိုပြောင်းလဲသွားလဲ print ထုတ်ကြည့်ပြီး dynamic_axes မပါဘဲ export လုပ်ရင် ဘာဖြစ်လာနိုင်လဲ စဉ်းစားပါ။

သတိလေးတစ်ချက်

dynamic_axes မသတ်မှတ်ဘဲ batch size 1 တစ်ခုတည်းနဲ့ export လုပ်ရင်၊ deployment code က batch size 8 ပါတဲ့ request တစ်ခု ပို့လိုက်တဲ့အခါ ONNX runtime က shape mismatch error ပေးနိုင်ပါတယ်။

model.eval() ခေါ်ဖို့ မေ့ကျန်ခဲ့ရင် dropout နဲ့ batchnorm layer တွေက training mode မှာ ကျန်ရှိနေပြီး graph ထဲမှာ randomness (သို့) batch-dependent statistics ကို ထည့်သွင်းသိမ်းဆည်းသွားနိုင်ပါတယ် — ရလဒ်အနေနဲ့ single-image inference အတွက် exported model ရဲ့ output က မူရင်း trained model နဲ့ ကွဲလွဲနေတတ်ပါတယ်။

PyTorch Docs — torch.onnx — Computer Vision