File size: 965 Bytes
b415c47 | 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 | #!/usr/bin/env python3
"""Minimal memory ONNX export. No validation, no onnxsim, tiny dummy shape.
If this fits in RAM, great. Otherwise the export must run on another machine."""
import gc, os, sys
import torch
from safetensors.torch import load_file
from traiNNer.archs.heart_arch import heart
ckpt, out, scale = sys.argv[1], sys.argv[2], int(sys.argv[3])
torch.set_num_threads(2)
torch.set_grad_enabled(False)
sd = load_file(ckpt, device="cpu")
m = heart(scale=scale, ape=False, use_checkpoint=False).eval()
m.load_state_dict(sd, strict=True)
del sd; gc.collect()
# tiny dummy to minimize activation memory during trace
x = torch.randn(1, 3, 64, 64)
torch.onnx.export(
m, x, out,
input_names=["input"], output_names=["output"],
dynamic_axes={"input": {0:"batch",2:"height",3:"width"}, "output": {0:"batch",2:"height_out",3:"width_out"}},
opset_version=17, do_constant_folding=True, dynamo=False,
)
del m, x; gc.collect()
print(f"OK -> {out}")
|