File size: 5,480 Bytes
5a8f833 | 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 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 | """Export NERVE (canonical, configurable size) or SPAN to dynamic-shape ONNX and verify.
- fp32 export with dynamic H/W on the lq/sr tensors
- onnx.checker
- ONNX Runtime CPU inference vs PyTorch CPU (max abs diff) at several shapes:
48x48, 323x711, 720x1280.
Random weights prove the graph is export-clean; pass a trained EMA checkpoint
(--checkpoint) to validate an actual trained model. Canonical export is
opset 20 (torch 2.14's dynamo exporter has an opset floor of 18; an opset-17
request still emits 18).
"""
from __future__ import annotations
import argparse
import json
import sys
from pathlib import Path
import numpy as np
import torch
from torch.export import Dim
from traiNNer.archs.nerve_arch import nerve
ROOT = Path(__file__).resolve().parents[2]
OUT_ROOT = ROOT / "experiments" / "nerve" / "export"
SHAPES: list[tuple[int, int]] = [(48, 48), (323, 711), (720, 1280)]
def load_checkpoint(model: torch.nn.Module, path: Path) -> None:
if path.suffix == ".safetensors":
try:
from safetensors.torch import load_file
sd = load_file(str(path), device="cpu")
except Exception: # noqa: BLE001 - legacy pickle with .safetensors name
import io
sd = torch.load(
io.BytesIO(path.read_bytes()), map_location="cpu", weights_only=True
)
else:
sd = torch.load(path, map_location="cpu", weights_only=True)
missing, unexpected = model.load_state_dict(sd, strict=True)
if missing or unexpected:
raise RuntimeError(
f"checkpoint mismatch: missing={missing} unexpected={unexpected}"
)
def build(arch: str, scale: int, dim: int, n_blocks: int) -> torch.nn.Module:
if arch == "span":
from traiNNer.archs.span_arch import span
return span(scale=scale)
return nerve(scale=scale, dim=dim, n_blocks=n_blocks)
def export_onnx(
model: torch.nn.Module,
scale: int,
opset: int,
out_path: Path,
) -> None:
model.eval()
x = torch.randn(1, 3, 64, 64)
dynamic_shapes = {"x": {2: Dim("H"), 3: Dim("W")}}
torch.onnx.export(
model,
(x,),
str(out_path),
opset_version=opset,
dynamic_shapes=dynamic_shapes,
)
def verify_ort(model: torch.nn.Module, onnx_path: Path, scale: int, tol: float) -> dict:
import onnx
import onnxruntime as ort
onnx.checker.check_model(onnx_path)
sess = ort.InferenceSession(str(onnx_path), providers=["CPUExecutionProvider"])
iname = sess.get_inputs()[0].name
model.eval()
results: dict[str, float] = {}
with torch.no_grad():
for h, w in SHAPES:
x = torch.randn(1, 3, h, w)
ref = model(x)
ort_out = sess.run(None, {iname: x.numpy()})[0]
diff = float(np.abs(ref.numpy() - ort_out).max())
results[f"{h}x{w}"] = round(diff, 6)
return results
def main() -> int:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--arch", default="nerve", choices=["nerve", "span"])
parser.add_argument("--dim", type=int, default=48)
parser.add_argument("--n-blocks", type=int, default=16)
parser.add_argument(
"--checkpoint", default=None, help="trained EMA safetensors to load"
)
parser.add_argument(
"--scale", type=int, default=4, help="scale when a checkpoint is given"
)
parser.add_argument("--scales", nargs="+", type=int, default=[4, 2])
parser.add_argument("--opsets", nargs="+", type=int, default=[20])
parser.add_argument("--tol", type=float, default=1e-3)
parser.add_argument("--out-root", default=str(OUT_ROOT))
args = parser.parse_args()
scales = [args.scale] if args.checkpoint else args.scales
overall_rc = 0
summary: dict = {}
for scale in scales:
model = build(args.arch, scale, args.dim, args.n_blocks)
model.to("cpu")
model.eval()
if args.checkpoint:
load_checkpoint(model, Path(args.checkpoint))
params = sum(p.numel() for p in model.parameters())
print(f"arch {args.arch} scale {scale} params={params:,}", flush=True)
scale_rc = 0
for opset in args.opsets:
tag = f"{args.arch}_d{args.dim}_b{args.n_blocks}_s{scale}"
out_dir = Path(args.out_root) / tag
out_dir.mkdir(parents=True, exist_ok=True)
out_path = out_dir / f"nerve_dynamic_op{opset}.onnx"
export_onnx(model, scale, opset, out_path)
print(f"exported {out_path.name}", flush=True)
try:
diffs = verify_ort(model, out_path, scale, args.tol)
except Exception as e: # noqa: BLE001
print(f"opset {opset} ORT verify FAILED: {e}", flush=True)
scale_rc = 1
continue
worst = max(diffs.values()) if diffs else float("inf")
ok = worst <= args.tol
print(
f"opset {opset} max|diff| per shape: {diffs} -> {'OK' if ok else 'FAIL'}",
flush=True,
)
scale_rc |= 0 if ok else 1
summary[f"arch{args.arch}_s{scale}_op{opset}"] = {
"max_diff": diffs,
"ok": ok,
}
overall_rc |= scale_rc
(Path(args.out_root) / "export_summary.json").write_text(
json.dumps(summary, indent=2)
)
return overall_rc
if __name__ == "__main__":
sys.exit(main())
|