File size: 1,358 Bytes
8f383b6 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 | # E-OS v5 — ONNX-Export (opset 17, App-Aufloesung 256x544 -> 512x1088) + Verify
import os
import sys
import numpy as np
import torch
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from os5_model import FrameGeneratorV5 # noqa: E402
CKPT = "/home/z/my-project/eos-v5/train/v5_ckpt.pt"
OUTD = "/home/z/my-project/eos-v5/models"
os.makedirs(OUTD, exist_ok=True)
model = FrameGeneratorV5()
ck = torch.load(CKPT, map_location="cpu", weights_only=False)
model.load_state_dict(ck["model"])
model.eval()
print("step:", ck["step"], "hist:", ck["hist"][-3:])
hint = torch.rand(1, 3, 544, 256)
state = torch.rand(1, 32)
with torch.no_grad():
ref = model(hint, state)
onnx_path = os.path.join(OUTD, "os_glass_v5.onnx")
torch.onnx.export(
model, (hint, state), onnx_path,
input_names=["hint", "state"], output_names=["frame"],
opset_version=17, dynamo=False,
do_constant_folding=True,
)
print("exported:", onnx_path, os.path.getsize(onnx_path) // 1024, "KB")
import onnxruntime as ort
so = ort.SessionOptions()
sess = ort.InferenceSession(onnx_path, so, providers=["CPUExecutionProvider"])
outs = sess.run(None, {"hint": hint.numpy(), "state": state.numpy()})
diff = np.abs(outs[0] - ref.numpy()).max()
print("ort output shape:", outs[0].shape, "max diff vs torch:", diff)
assert diff < 1e-4, "ONNX mismatch"
print("ONNX_OK")
|