[参考译文] PROCESSOR-SDK-J722S：关于将 onnx 文件转换为 TIDL 格式的问题

import torch
import onnx
import onnxruntime as rt
import subprocess
import os

# 定義參數
onnx_model_path = "model.onnx"
simplified_onnx_model_path = "model_simplified.onnx"
tidl_artifacts_folder = "./model-artifacts-dir/"
tidl_tools_path = "./app/ken/tda4/TIDL10.00.06.00/edgeai-tidl-tools/tidl_tools/"

# Step 1: 導出 ONNX 模型 (以 PyTorch 為例)
def export_pytorch_model_to_onnx(model, input_shape):
    dummy_input = torch.randn(*input_shape)
    torch.onnx.export(model, dummy_input, onnx_model_path, opset_version=11)
    print(f"ONNX model exported to {onnx_model_path}")

# Step 2: 檢查模型有效性並推斷形狀
def check_and_infer_shape(onnx_model_path):
    model = onnx.load(onnx_model_path)
    inferred_model = onnx.shape_inference.infer_shapes(model)
    onnx.save(inferred_model, "inferred_model.onnx")
    print("Shape inference complete and saved to inferred_model.onnx")

# Step 3: 簡化 ONNX 模型
def simplify_onnx_model(onnx_model_path, simplified_model_path):
    subprocess.run(["python3", "-m", "onnxsim", onnx_model_path, simplified_model_path])
    print(f"Simplified ONNX model saved to {simplified_model_path}")

# Step 4: 設定 ONNX Runtime TIDL 編譯會話
def setup_tidl_session(simplified_model_path):
    # 設置編譯參數
    options = {
        "artifacts_folder": tidl_artifacts_folder,
        "tidl_tools_path": tidl_tools_path
    }
    so = rt.SessionOptions()
    ep_list = ['TIDLCompilationProvider', 'CPUExecutionProvider']
    
    # 設置編譯會話，使用 TIDLCompilationProvider 進行編譯
    sess = rt.InferenceSession(simplified_model_path, providers=ep_list, provider_options=[options, {}], sess_options=so)
    
    print("TIDL model compilation session created.")
    return sess

# Step 5: 驗證模型的輸出
def validate_model_output(sess, input_data):
    # 確保 TIDL 編譯模型可以產生正確輸出
    input_name = sess.get_inputs()[0].name
    output_name = sess.get_outputs()[0].name
    
    output = sess.run([output_name], {input_name: input_data})
    print(f"Model output: {output}")

# 主流程
if __name__ == "__main__":
    # 假設這是 PyTorch 模型和輸入維度
    model = torch.nn.Sequential(
        torch.nn.Conv2d(3, 16, 3, stride=2, padding=1),
        torch.nn.ReLU(),
        torch.nn.Flatten(),
        torch.nn.Linear(16 * 112 * 112, 10)
    )
    input_shape = (1, 3, 224, 224)

    # 導出、檢查、簡化、編譯和驗證模型
    export_pytorch_model_to_onnx(model, input_shape)
    check_and_infer_shape(onnx_model_path)
    simplify_onnx_model(onnx_model_path, simplified_onnx_model_path)
    session = setup_tidl_session(simplified_onnx_model_path)
    
    # 使用隨機數據進行驗證
    input_data = torch.randn(*input_shape).numpy()
    validate_model_output(session, input_data)
    print("Available providers:", rt.get_available_providers())