| """A trained direct predictor as an ONNX file under the benchmark's contract: python export.py NAME OUT.onnx |
| |
| Reads ml/runs/NAME.pt and writes a graph from boundary_r_cos, boundary_z_sin (5, 9), requirement_metrics (4), |
| n_field_periods and n_coils_per_half_period to coil_x_n, coil_y_n, coil_z_n (K, 15) and currents (K) in amperes, |
| with the metadata key supported_tracks set to fixed_shape. The scalings of the inputs and outputs are constants of |
| the graph. Prints the largest difference between the graph's outputs in onnxruntime and the network's in PyTorch. |
| """ |
| from __future__ import annotations |
|
|
| import sys |
| from pathlib import Path |
|
|
| import numpy as np |
| import onnx |
| import onnxruntime as ort |
| import torch |
| from torch import nn |
|
|
| from train import Net |
|
|
| HERE = Path(__file__).resolve().parent |
|
|
|
|
| class Wrapped(nn.Module): |
| def __init__(self, ck): |
| super().__init__() |
| self.k = int(ck["k"]) |
| self.net = Net(int(ck["varies"].sum()), self.k * 45 + self.k, int(ck["width"]), int(ck["depth"])) |
| self.net.load_state_dict(ck["net"]) |
| self.register_buffer("pick", torch.nonzero(ck["varies"]).reshape(-1)) |
| for name in ("x_mean", "x_std", "y_mean", "y_std"): |
| self.register_buffer(name, ck[name].float()) |
|
|
| def forward(self, boundary_r_cos, boundary_z_sin, requirement_metrics, n_field_periods, n_coils_per_half_period): |
| x = torch.cat([boundary_r_cos.reshape(45), boundary_z_sin.reshape(45), requirement_metrics.log()]) |
| x = (x.index_select(0, self.pick) - self.x_mean) / self.x_std |
| y = self.net(x.reshape(1, -1)).reshape(-1) * self.y_std + self.y_mean |
| |
| y = y + 0.0 * (n_field_periods + n_coils_per_half_period).to(y.dtype) |
| coef = y[:self.k * 45].reshape(self.k, 3, 15) |
| return coef[:, 0], coef[:, 1], coef[:, 2], y[self.k * 45:] * 1e6 |
|
|
|
|
| def main(name, out): |
| ck = torch.load(HERE / "runs" / f"{name}.pt", map_location="cpu") |
| model = Wrapped(ck).eval() |
| feed = (torch.randn(5, 9) * 0.05, torch.randn(5, 9) * 0.05, torch.tensor([0.4, 0.8, 3.0, 0.05]), torch.tensor(3), torch.tensor(5)) |
| feed[0][0, 4] = 1.0 |
| names_in = ["boundary_r_cos", "boundary_z_sin", "requirement_metrics", "n_field_periods", "n_coils_per_half_period"] |
| names_out = ["coil_x_n", "coil_y_n", "coil_z_n", "currents"] |
| torch.onnx.export(model, feed, out, input_names=names_in, output_names=names_out, opset_version=17, dynamo=False) |
| proto = onnx.load(out) |
| proto.metadata_props.append(onnx.StringStringEntryProto(key="supported_tracks", value="fixed_shape")) |
| onnx.save(proto, out) |
| onnx.checker.check_model(out) |
| sess = ort.InferenceSession(out, providers=["CPUExecutionProvider"]) |
| got = sess.run(None, {n: (v.numpy().astype(np.int64) if v.dtype == torch.int64 else v.numpy()) for n, v in zip(names_in, feed)}) |
| with torch.no_grad(): |
| want = model(*feed) |
| worst = max(float(np.abs(g - w.numpy()).max() / max(1.0, float(w.abs().max()))) for g, w in zip(got, want)) |
| print(f"wrote {out}: {Path(out).stat().st_size / 2 ** 20:.1f} MiB, dev feasible in training {ck.get('dev_feasible')}, " |
| f"largest relative difference from PyTorch {worst:.2e}, inputs {[i.name for i in sess.get_inputs()]}") |
|
|
|
|
| if __name__ == "__main__": |
| main(*sys.argv[1:3]) |
|
|
