Upload own-solver/neurogolf_solver/onnx_helpers.py

own-solver/neurogolf_solver/onnx_helpers.py

@@ -0,0 +1,67 @@
+#!/usr/bin/env python3
+"""ONNX model building helper functions (opset 17)."""
+
+import numpy as np
+from onnx import helper, TensorProto, numpy_helper
+from .constants import DT, IR, GRID_SHAPE, INT64_MIN, GH, GW
+from .config import make_opset
+
+
+def _make_int64_init(name, values):
+    """Create int64 initializer."""
+    return numpy_helper.from_array(np.array(values, dtype=np.int64), name)
+
+
+def _build_pad_node(input_name, output_name, pad_h, pad_w, inits, suffix=''):
+    """Pad with tensor-based pads input (opset 11+)."""
+    pads_name = f'pads{suffix}'
+    cv_name = f'pad_cv{suffix}'
+    pads_arr = np.array([0, 0, 0, 0, 0, 0, pad_h, pad_w], dtype=np.int64)
+    inits.append(numpy_helper.from_array(pads_arr, pads_name))
+    inits.append(numpy_helper.from_array(np.array(0.0, dtype=np.float32), cv_name))
+    return helper.make_node('Pad', [input_name, pads_name, cv_name], [output_name], mode='constant')
+
+
+def _build_slice_crop(input_name, output_name, IH, IW, inits, suffix=''):
+    """Slice to crop [1,10,30,30] to [1,10,IH,IW]."""
+    st_name = f'crop_st{suffix}'
+    en_name = f'crop_en{suffix}'
+    inits.append(_make_int64_init(st_name, [0, 0, 0, 0]))
+    inits.append(_make_int64_init(en_name, [1, 10, IH, IW]))
+    return helper.make_node('Slice', [input_name, st_name, en_name], [output_name])
+
+
+def _build_slice_reverse(input_name, output_name, axis, dim_size, inits, suffix=''):
+    """Slice(step=-1) to reverse one axis. Zero MACs."""
+    st_name = f'rev_st{suffix}'
+    en_name = f'rev_en{suffix}'
+    ax_name = f'rev_ax{suffix}'
+    sp_name = f'rev_sp{suffix}'
+    inits.append(_make_int64_init(st_name, [dim_size - 1]))
+    inits.append(_make_int64_init(en_name, [INT64_MIN]))
+    inits.append(_make_int64_init(ax_name, [axis]))
+    inits.append(_make_int64_init(sp_name, [-1]))
+    return helper.make_node('Slice', [input_name, st_name, en_name, ax_name, sp_name], [output_name])
+
+
+def _build_reducesum(input_name, output_name, axes_list, inits, suffix=''):
+    """ReduceSum with axes as tensor input (opset 13+). keepdims=1."""
+    axes_name = f'rs_axes{suffix}'
+    inits.append(_make_int64_init(axes_name, axes_list))
+    return helper.make_node('ReduceSum', [input_name, axes_name], [output_name], keepdims=1)
+
+
+def mk(nodes, inits=None, opset_version=17):
+    """Create ONNX model from nodes and initializers."""
+    x = helper.make_tensor_value_info("input", DT, GRID_SHAPE)
+    y = helper.make_tensor_value_info("output", DT, GRID_SHAPE)
+    g = helper.make_graph(nodes, "g", [x], [y], initializer=inits or [])
+    return helper.make_model(g, ir_version=IR, opset_imports=make_opset(opset_version))
+
+
+def add_onehot_block(nodes, inits, am_name, oh_name):
+    """Add ArgMax one-hot conversion block."""
+    classes = np.arange(10, dtype=np.int64).reshape(1, 10, 1, 1)
+    inits.append(numpy_helper.from_array(classes, 'classes'))
+    nodes.append(helper.make_node('Equal', [am_name, 'classes'], ['eq']))
+    nodes.append(helper.make_node('Cast', ['eq'], [oh_name], to=TensorProto.FLOAT))