50 lines
1.7 KiB
Python
50 lines
1.7 KiB
Python
"""
|
|
export_onnx.py — export a trained U-Net to ONNX, auto-detecting channels (v1=2, v2=3).
|
|
|
|
python src\\export_onnx.py --ckpt checkpoints\\best.pt --out web\\radio_unet_v1.onnx
|
|
python src\\export_onnx.py --ckpt checkpoints_v2\\best.pt --out web\\radio_unet_v2.onnx
|
|
"""
|
|
import argparse
|
|
import numpy as np
|
|
import torch
|
|
|
|
from model import UNet
|
|
from dataset import RadioMapSeerDataset
|
|
|
|
|
|
def main():
|
|
p = argparse.ArgumentParser()
|
|
p.add_argument("--ckpt", default="checkpoints/best.pt")
|
|
p.add_argument("--out", default="web/radio_unet.onnx")
|
|
args = p.parse_args()
|
|
|
|
ckpt = torch.load(args.ckpt, map_location="cpu", weights_only=False)
|
|
cargs = ckpt.get("args", {})
|
|
base = cargs.get("base", 64)
|
|
add_dist = bool(cargs.get("distance", False))
|
|
in_ch = 3 if add_dist else 2
|
|
|
|
model = UNet(in_channels=in_ch, out_channels=1, base=base)
|
|
model.load_state_dict(ckpt["model_state"])
|
|
model.eval()
|
|
print(f"Loaded {args.ckpt} | in_channels={in_ch} | distance={add_dist}")
|
|
|
|
dummy = torch.randn(1, in_ch, 256, 256)
|
|
torch.onnx.export(model, dummy, args.out,
|
|
input_names=["input"], output_names=["coverage"],
|
|
opset_version=17)
|
|
print(f"Exported {args.out}")
|
|
|
|
import onnxruntime as ort
|
|
ds = RadioMapSeerDataset(split="test", img_size=256, add_distance=add_dist)
|
|
x, _ = ds[0]
|
|
xb = x.unsqueeze(0).numpy().astype(np.float32)
|
|
with torch.no_grad():
|
|
t_out = model(x.unsqueeze(0)).numpy()
|
|
sess = ort.InferenceSession(args.out, providers=["CPUExecutionProvider"])
|
|
o_out = sess.run(["coverage"], {"input": xb})[0]
|
|
print(f"Max |PyTorch - ONNX|: {np.max(np.abs(t_out - o_out)):.2e}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main() |