# ============================================================ # 0. Install dependencies (optional, uncomment if needed) # ============================================================ # !pip -q install onnx onnxruntime onnxsim torch huggingface_hub import os import json import math import zipfile import numpy as np import onnx from onnx import helper, numpy_helper, TensorProto from onnxruntime.quantization import ( quantize_static, quantize_dynamic, QuantFormat, QuantType, CalibrationDataReader, ) import subprocess import shutil # Safely import onnxsim with fallback try: from onnxsim import simplify HAS_ONNXSIM = True except ImportError: HAS_ONNXSIM = False simplify = None import torch import torch.nn as nn import torch.nn.functional as F # ============================================================ # 1. Build v4 + reducer ONNX graph # ============================================================ def build_task062_v4_reducer(): nodes, inits, vis = [], [], [] counter = [0] def nm(): counter[0] += 1 return f"t{counter[0]}" def const(name, val, dtype='f'): arr = np.array(val, dtype=np.float32 if dtype == 'f' else np.int64) inits.append(numpy_helper.from_array(arr, name)) def vi(name, shape, dt=TensorProto.FLOAT): vis.append(helper.make_tensor_value_info(name, dt, shape)) def nd(op, ins, out_shape, **kwargs): if isinstance(out_shape, tuple): shape, dt = out_shape else: shape, dt = out_shape, TensorProto.FLOAT out = nm() vi(out, shape, dt) nodes.append(helper.make_node(op, ins, [out], **kwargs)) return out # ----------------------- # Constants # ----------------------- const('c_half', [0.5]) const('c_one', [1.0]) const('c_zero', [0.0]) const('c_two', [2.0]) const('c_big', [100.0]) const('c_neg_big', [-100.0]) const('sl_start', [0, 0, 0, 0], 'i') const('sl_end', [1, 10, 10, 10], 'i') const('sl_axes', [0, 1, 2, 3], 'i') const('axes1', [1], 'i') const('axes2', [2], 'i') const('axes3', [3], 'i') const('axes23', [2, 3], 'i') const('row_grid', np.arange(10, dtype=np.float32).reshape(1, 1, 10, 1)) const('col_grid', np.arange(10, dtype=np.float32).reshape(1, 1, 1, 10)) const('row_1d', np.arange(10, dtype=np.float32).reshape(1, 1, 10, 1)) const('col_1d', np.arange(10, dtype=np.float32).reshape(1, 1, 1, 10)) const('shape_1', [1], 'i') const('shape_1_10_1_1', [1, 10, 1, 1], 'i') const('depth_10', [10.0]) const('oh_vals', [0.0, 1.0]) const('pad_10_to_30', [0, 0, 0, 0, 0, 0, 20, 20], 'i') const('pad_val_zero', [0.0]) const('nobg_no2_mask', np.array([[[[0, 1, 0, 1, 1, 1, 1, 1, 1, 1]]]], dtype=np.float32).reshape(1, 10, 1, 1)) const('ch2_start', [0, 2, 0, 0], 'i') const('ch2_end', [1, 3, 10, 10], 'i') const('idx_3', [3], 'i') const('ones_1_1_1_10', np.ones((1, 1, 1, 10), dtype=np.float32)) const('ones_1_1_10_1', np.ones((1, 1, 10, 1), dtype=np.float32)) const('c_0f', [0.0]) const('c_9_0', [9.0]) # ----------------------- # Input/Output # ----------------------- x = helper.make_tensor_value_info('input', TensorProto.FLOAT, [1, 10, 30, 30]) y = helper.make_tensor_value_info('output', TensorProto.FLOAT, [1, 10, 30, 30]) # ----------------------- # Crop # ----------------------- inp10 = nd('Slice', ['input', 'sl_start', 'sl_end', 'sl_axes'], [1, 10, 10, 10]) # ----------------------- # Main color # ----------------------- ch_sums = nd('ReduceSum', [inp10, 'axes23'], [1, 10, 1, 1], keepdims=1) ch_sums_m = nd('Mul', [ch_sums, 'nobg_no2_mask'], [1, 10, 1, 1]) mc_idx = nd('ArgMax', [ch_sums_m], ([1, 1, 1, 1], TensorProto.INT64), axis=1, keepdims=1) mc_1d = nd('Reshape', [mc_idx, 'shape_1'], ([1], TensorProto.INT64)) mc_oh = nd('OneHot', [mc_1d, 'depth_10', 'oh_vals'], [1, 10], axis=1) mc_oh_4d = nd('Reshape', [mc_oh, 'shape_1_10_1_1'], [1, 10, 1, 1]) main_sel = nd('Mul', [inp10, mc_oh_4d], [1, 10, 10, 10]) main_mask = nd('ReduceSum', [main_sel, 'axes1'], [1, 1, 10, 10], keepdims=1) axis_mask = nd('Slice', [inp10, 'ch2_start', 'ch2_end', 'sl_axes'], [1, 1, 10, 10]) # ----------------------- # Projections # ----------------------- main_rp = nd('ReduceSum', [main_mask, 'axes3'], [1, 1, 10, 1], keepdims=1) main_cp = nd('ReduceSum', [main_mask, 'axes2'], [1, 1, 1, 10], keepdims=1) axis_rp = nd('ReduceSum', [axis_mask, 'axes3'], [1, 1, 10, 1], keepdims=1) axis_cp = nd('ReduceSum', [axis_mask, 'axes2'], [1, 1, 1, 10], keepdims=1) # ----------------------- # Centroids # ----------------------- main_total = nd('ReduceSum', [main_rp, 'axes2'], [1, 1, 1, 1], keepdims=1) axis_total = nd('ReduceSum', [axis_rp, 'axes2'], [1, 1, 1, 1], keepdims=1) main_rp_row = nd('Mul', [main_rp, 'row_1d'], [1, 1, 10, 1]) main_rp_row_sum = nd('ReduceSum', [main_rp_row, 'axes2'], [1, 1, 1, 1], keepdims=1) main_cr = nd('Div', [main_rp_row_sum, main_total], [1, 1, 1, 1]) main_cp_col = nd('Mul', [main_cp, 'col_1d'], [1, 1, 1, 10]) main_cp_col_sum = nd('ReduceSum', [main_cp_col, 'axes3'], [1, 1, 1, 1], keepdims=1) main_cc = nd('Div', [main_cp_col_sum, main_total], [1, 1, 1, 1]) axis_rp_row = nd('Mul', [axis_rp, 'row_1d'], [1, 1, 10, 1]) axis_rp_row_sum = nd('ReduceSum', [axis_rp_row, 'axes2'], [1, 1, 1, 1], keepdims=1) axis_cr = nd('Div', [axis_rp_row_sum, axis_total], [1, 1, 1, 1]) axis_cp_col = nd('Mul', [axis_cp, 'col_1d'], [1, 1, 1, 10]) axis_cp_col_sum = nd('ReduceSum', [axis_cp_col, 'axes3'], [1, 1, 1, 1], keepdims=1) axis_cc = nd('Div', [axis_cp_col_sum, axis_total], [1, 1, 1, 1]) # ----------------------- # Orientation # ----------------------- dr = nd('Sub', [axis_cr, main_cr], [1, 1, 1, 1]) dc = nd('Sub', [axis_cc, main_cc], [1, 1, 1, 1]) abs_dr = nd('Abs', [dr], [1, 1, 1, 1]) abs_dc = nd('Abs', [dc], [1, 1, 1, 1]) abs_dc_minus_half = nd('Sub', [abs_dc, 'c_half'], [1, 1, 1, 1]) is_horiz_b = nd('Greater', [abs_dr, abs_dc_minus_half], ([1, 1, 1, 1], TensorProto.BOOL)) is_horiz = nd('Cast', [is_horiz_b], [1, 1, 1, 1], to=1) is_vert = nd('Sub', ['c_one', is_horiz], [1, 1, 1, 1]) # ----------------------- # Edge reducers # ----------------------- has_mr_b = nd('Greater', [main_rp, 'c_half'], ([1, 1, 10, 1], TensorProto.BOOL)) has_mr = nd('Cast', [has_mr_b], [1, 1, 10, 1], to=1) no_mr = nd('Sub', ['c_one', has_mr], [1, 1, 10, 1]) mr_for_max = nd('Add', [ nd('Mul', ['row_1d', has_mr], [1, 1, 10, 1]), nd('Mul', ['c_neg_big', no_mr], [1, 1, 10, 1]) ], [1, 1, 10, 1]) max_mr = nd('ReduceMax', [mr_for_max, 'axes2'], [1, 1, 1, 1], keepdims=1) mr_for_min = nd('Add', [ nd('Mul', ['row_1d', has_mr], [1, 1, 10, 1]), nd('Mul', ['c_big', no_mr], [1, 1, 10, 1]) ], [1, 1, 10, 1]) min_mr = nd('ReduceMin', [mr_for_min, 'axes2'], [1, 1, 1, 1], keepdims=1) has_mc_b = nd('Greater', [main_cp, 'c_half'], ([1, 1, 1, 10], TensorProto.BOOL)) has_mc = nd('Cast', [has_mc_b], [1, 1, 1, 10], to=1) no_mc = nd('Sub', ['c_one', has_mc], [1, 1, 1, 10]) mc_for_max = nd('Add', [ nd('Mul', ['col_1d', has_mc], [1, 1, 1, 10]), nd('Mul', ['c_neg_big', no_mc], [1, 1, 1, 10]) ], [1, 1, 1, 10]) max_mc = nd('ReduceMax', [mc_for_max, 'axes3'], [1, 1, 1, 1], keepdims=1) mc_for_min = nd('Add', [ nd('Mul', ['col_1d', has_mc], [1, 1, 1, 10]), nd('Mul', ['c_big', no_mc], [1, 1, 1, 10]) ], [1, 1, 1, 10]) min_mc = nd('ReduceMin', [mc_for_min, 'axes3'], [1, 1, 1, 1], keepdims=1) has_ar_b = nd('Greater', [axis_rp, 'c_half'], ([1, 1, 10, 1], TensorProto.BOOL)) has_ar = nd('Cast', [has_ar_b], [1, 1, 10, 1], to=1) no_ar = nd('Sub', ['c_one', has_ar], [1, 1, 10, 1]) ar_for_max = nd('Add', [ nd('Mul', ['row_1d', has_ar], [1, 1, 10, 1]), nd('Mul', ['c_neg_big', no_ar], [1, 1, 10, 1]) ], [1, 1, 10, 1]) max_ar = nd('ReduceMax', [ar_for_max, 'axes2'], [1, 1, 1, 1], keepdims=1) ar_for_min = nd('Add', [ nd('Mul', ['row_1d', has_ar], [1, 1, 10, 1]), nd('Mul', ['c_big', no_ar], [1, 1, 10, 1]) ], [1, 1, 10, 1]) min_ar = nd('ReduceMin', [ar_for_min, 'axes2'], [1, 1, 1, 1], keepdims=1) has_ac_b = nd('Greater', [axis_cp, 'c_half'], ([1, 1, 1, 10], TensorProto.BOOL)) has_ac = nd('Cast', [has_ac_b], [1, 1, 1, 10], to=1) no_ac = nd('Sub', ['c_one', has_ac], [1, 1, 1, 10]) ac_for_max = nd('Add', [ nd('Mul', ['col_1d', has_ac], [1, 1, 1, 10]), nd('Mul', ['c_neg_big', no_ac], [1, 1, 1, 10]) ], [1, 1, 1, 10]) max_ac = nd('ReduceMax', [ac_for_max, 'axes3'], [1, 1, 1, 1], keepdims=1) ac_for_min = nd('Add', [ nd('Mul', ['col_1d', has_ac], [1, 1, 1, 10]), nd('Mul', ['c_big', no_ac], [1, 1, 1, 10]) ], [1, 1, 1, 10]) min_ac = nd('ReduceMin', [ac_for_min, 'axes3'], [1, 1, 1, 1], keepdims=1) # ----------------------- # Mirror positions # ----------------------- dr_pos_b = nd('Greater', [dr, 'c_zero'], ([1, 1, 1, 1], TensorProto.BOOL)) dr_pos = nd('Cast', [dr_pos_b], [1, 1, 1, 1], to=1) dr_neg = nd('Sub', ['c_one', dr_pos], [1, 1, 1, 1]) h_sum = nd('Add', [ nd('Add', [ nd('Mul', [max_mr, dr_pos], [1, 1, 1, 1]), nd('Mul', [min_mr, dr_neg], [1, 1, 1, 1]) ], [1, 1, 1, 1]), nd('Add', [ nd('Mul', [min_ar, dr_pos], [1, 1, 1, 1]), nd('Mul', [max_ar, dr_neg], [1, 1, 1, 1]) ], [1, 1, 1, 1]) ], [1, 1, 1, 1]) h_mirror = nd('Div', [h_sum, 'c_two'], [1, 1, 1, 1]) dc_pos_b = nd('Greater', [dc, 'c_zero'], ([1, 1, 1, 1], TensorProto.BOOL)) dc_pos = nd('Cast', [dc_pos_b], [1, 1, 1, 1], to=1) dc_neg = nd('Sub', ['c_one', dc_pos], [1, 1, 1, 1]) v_sum = nd('Add', [ nd('Add', [ nd('Mul', [max_mc, dc_pos], [1, 1, 1, 1]), nd('Mul', [min_mc, dc_neg], [1, 1, 1, 1]) ], [1, 1, 1, 1]), nd('Add', [ nd('Mul', [min_ac, dc_pos], [1, 1, 1, 1]), nd('Mul', [max_ac, dc_neg], [1, 1, 1, 1]) ], [1, 1, 1, 1]) ], [1, 1, 1, 1]) v_mirror = nd('Div', [v_sum, 'c_two'], [1, 1, 1, 1]) # ----------------------- # Reflection maps # ----------------------- h_mx2 = nd('Mul', [h_mirror, 'c_two'], [1, 1, 1, 1]) h_rr = nd('Floor', [nd('Add', [nd('Sub', [h_mx2, 'row_grid'], [1, 1, 10, 1]), 'c_half'], [1, 1, 10, 1])], [1, 1, 10, 1]) h_rc = nd('Clip', [h_rr, 'c_0f', 'c_9_0'], [1, 1, 10, 1]) h_idx = nd('Cast', [nd('Mul', [h_rc, 'ones_1_1_1_10'], [1, 1, 10, 10])], ([1, 1, 10, 10], TensorProto.INT64), to=7) h_refl = nd('GatherElements', [main_mask, h_idx], [1, 1, 10, 10], axis=2) v_mx2 = nd('Mul', [v_mirror, 'c_two'], [1, 1, 1, 1]) v_rc = nd('Floor', [nd('Add', [nd('Sub', [v_mx2, 'col_grid'], [1, 1, 1, 10]), 'c_half'], [1, 1, 1, 10])], [1, 1, 1, 10]) v_rclip = nd('Clip', [v_rc, 'c_0f', 'c_9_0'], [1, 1, 1, 10]) v_idx = nd('Cast', [nd('Mul', [v_rclip, 'ones_1_1_10_1'], [1, 1, 10, 10])], ([1, 1, 10, 10], TensorProto.INT64), to=7) v_refl = nd('GatherElements', [main_mask, v_idx], [1, 1, 10, 10], axis=3) # ----------------------- # Combine # ----------------------- refl_mask = nd('Add', [ nd('Mul', [is_horiz, h_refl], [1, 1, 10, 10]), nd('Mul', [is_vert, v_refl], [1, 1, 10, 10]) ], [1, 1, 10, 10]) combined = nd('Max', [main_mask, refl_mask], [1, 1, 10, 10]) # ----------------------- # Output # ----------------------- main_out = nd('Mul', [mc_oh_4d, combined], [1, 10, 10, 10]) bg_mask = nd('Sub', ['c_one', combined], [1, 1, 10, 10]) c3_oh = nd('OneHot', ['idx_3', 'depth_10', 'oh_vals'], [1, 10], axis=1) c3_oh_4d = nd('Reshape', [c3_oh, 'shape_1_10_1_1'], [1, 10, 1, 1]) bg_out = nd('Mul', [c3_oh_4d, bg_mask], [1, 10, 10, 10]) out_10 = nd('Add', [main_out, bg_out], [1, 10, 10, 10]) final = nd('Pad', [out_10, 'pad_10_to_30', 'pad_val_zero'], [1, 10, 30, 30]) nodes.append(helper.make_node('Identity', [final], ['output'])) graph = helper.make_graph(nodes, 'task062_v4_reducer', [x], [y], initializer=inits, value_info=vis) model = helper.make_model(graph, ir_version=10, opset_imports=[helper.make_opsetid('', 18)]) return model # ============================================================ # 2. Paths # ============================================================ OUT_DIR = "/kaggle/working/onnx_out" os.makedirs(OUT_DIR, exist_ok=True) RAW_ONNX = os.path.join(OUT_DIR, "task062_raw.onnx") SIM_ONNX = os.path.join(OUT_DIR, "task062_simplified.onnx") PREPROC_ONNX = os.path.join(OUT_DIR, "task062_preprocessed.onnx") INT8_ONNX = os.path.join(OUT_DIR, "task062_int8.onnx") DYN_ONNX = os.path.join(OUT_DIR, "task062_dynamic.onnx") # ============================================================ # 3. Build and export model # ============================================================ print("Building v4+reducer model...") model = build_task062_v4_reducer() print("Inferring shapes...") del model.graph.value_info[:] model = onnx.shape_inference.infer_shapes(model, strict_mode=True) onnx.save(model, RAW_ONNX) print("Saved raw ONNX:", RAW_ONNX, "size:", os.path.getsize(RAW_ONNX), "bytes") # ============================================================ # 4. Optional simplification # ============================================================ print("Simplifying (if onnxsim available)...") if HAS_ONNXSIM: onnx_model = onnx.load(RAW_ONNX) model_simp, ok = simplify(onnx_model) assert ok, "ONNX simplification failed" onnx.save(model_simp, SIM_ONNX) onnx_for_preproc = SIM_ONNX print("Saved simplified ONNX:", SIM_ONNX, "size:", os.path.getsize(SIM_ONNX), "bytes") else: onnx.save(model, SIM_ONNX) onnx_for_preproc = RAW_ONNX print("onnxsim not available; using raw ONNX for pre-processing") # ============================================================ # 5. ONNX Runtime pre-processing (always safe) # ============================================================ print("Running ONNX Runtime pre-processing (optional)...") cmd = [ "python", "-m", "onnxruntime.quantization.preprocess", "--input", onnx_for_preproc, "--output", PREPROC_ONNX, ] result = subprocess.run(cmd, capture_output=True, text=True) if result.returncode != 0: print("Pre-processing failed (symbolic shape inference issue).") print("Falling back to simplified model as preprocessed input.") shutil.copy(onnx_for_preproc, PREPROC_ONNX) print("Preprocessing skipped: using simplified ONNX directly.") else: print("Pre-processing succeeded.") print(result.stdout) print("Preprocessed ONNX:", PREPROC_ONNX, "size:", os.path.getsize(PREPROC_ONNX), "bytes") # ============================================================ # 6. Calibration data reader # ============================================================ class MyCalibrationDataReader(CalibrationDataReader): def __init__(self, batches): super().__init__() self.batches = iter(batches) def get_next(self): try: x = next(self.batches) return {"input": x.numpy()} except StopIteration: return None # Replace with real calibration samples if available calibration_batches = [torch.randn(1, 10, 30, 30) for _ in range(64)] reader = MyCalibrationDataReader(calibration_batches) # ============================================================ # 7. Static and dynamic quantization # ============================================================ quant_input = PREPROC_ONNX print("Static INT8 quantization...") quantize_static( quant_input, INT8_ONNX, reader, quant_format=QuantFormat.QDQ, activation_type=QuantType.QInt8, weight_type=QuantType.QInt8, ) print("Saved INT8 ONNX:", INT8_ONNX, "size:", os.path.getsize(INT8_ONNX), "bytes") print("Dynamic quantization fallback...") quantize_dynamic( quant_input, DYN_ONNX, weight_type=QuantType.QInt8, ) print("Saved dynamic ONNX:", DYN_ONNX, "size:", os.path.getsize(DYN_ONNX), "bytes") # ============================================================ # 8. Size report # ============================================================ report = [] for f in [RAW_ONNX, SIM_ONNX, PREPROC_ONNX, INT8_ONNX, DYN_ONNX]: report.append({ "file": os.path.basename(f), "size_mb": round(os.path.getsize(f) / (1024 * 1024), 4) }) print("\nSize report:") print(json.dumps(report, indent=2))