| |
| 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 |
|
|
| 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") |
|
|
