Upload neurogolf_utils.py
Browse files- neurogolf_utils.py +559 -0
neurogolf_utils.py
ADDED
|
@@ -0,0 +1,559 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Copyright 2026 Google LLC
|
| 2 |
+
#
|
| 3 |
+
# Licensed under the Apache License, Version 2.0 (the "License");
|
| 4 |
+
# you may not use this file except in compliance with the License.
|
| 5 |
+
# You may obtain a copy of the License at
|
| 6 |
+
#
|
| 7 |
+
# https://www.apache.org/licenses/LICENSE-2.0
|
| 8 |
+
#
|
| 9 |
+
# Unless required by applicable law or agreed to in writing, software
|
| 10 |
+
# distributed under the License is distributed on an "AS IS" BASIS,
|
| 11 |
+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
| 12 |
+
# See the License for the specific language governing permissions and
|
| 13 |
+
# limitations under the License.
|
| 14 |
+
|
| 15 |
+
"""
|
| 16 |
+
Module containing utilities for the IJCAI-ECAI 2026 NeuroGolf Championship.
|
| 17 |
+
|
| 18 |
+
Version History:
|
| 19 |
+
* 2026-05-14:
|
| 20 |
+
* Puts back the _EXCLUDED_OP_TYPES check (thank you @cdeotte and @pavelsavchenkov!)
|
| 21 |
+
* Applies stronger sanitization to tensor & node names (thank you @linkinpony!)
|
| 22 |
+
* Rejects duplicate graph.value_info entries with the same tensor name (thank you @pavelsavchenkov!)
|
| 23 |
+
* 2026-05-06:
|
| 24 |
+
* Scalar parameters are now penalized with unit cost.
|
| 25 |
+
* Each tensor's memory footprint is set to the maximum size across all runs.
|
| 26 |
+
* Duplicate node names no longer create parameter undercount.
|
| 27 |
+
* Tensor names containing ONNX's special "kernel_time" string are disallowed.
|
| 28 |
+
* Runtime trace file prefixes are specified to prevent profile clobbering.
|
| 29 |
+
* Multi-input / multi-output graphs disallowed.
|
| 30 |
+
* 2026-05-04:
|
| 31 |
+
* Sequences and nonpositive tensor dimensions are disallowed.
|
| 32 |
+
* Accurate shape information derived from the ONNX Runtime Profiler.
|
| 33 |
+
* MACs no longer contribute to the objective criterion.
|
| 34 |
+
* 2026-05-04:
|
| 35 |
+
* Sequences and nonpositive tensor dimensions are disallowed.
|
| 36 |
+
* Accurate shape information derived from the ONNX Runtime Profiler.
|
| 37 |
+
* MACs no longer contribute to the objective criterion.
|
| 38 |
+
* 2026-04-30:
|
| 39 |
+
* Compress operators have been banned.
|
| 40 |
+
* Name collision between tensors and initializers are disallowed.
|
| 41 |
+
* Functions / custom domains / subgraphs are disallowed.
|
| 42 |
+
* Zero-cost networks now yield a full 25 points.
|
| 43 |
+
* 2026-04-28:
|
| 44 |
+
* Constant folding enabled to address the undercounting of parameters.
|
| 45 |
+
* Our "statically-defined shapes" constaint is now strictly enforced.
|
| 46 |
+
* Memory footprint calculation is now a sum of static shape sizes.
|
| 47 |
+
* Nodes with negative parameter counts or MACs are disallowed.
|
| 48 |
+
* 2026-04-21:
|
| 49 |
+
* Tests with grids larger than 30x30 are ignored.
|
| 50 |
+
* Nodes with negative memory values are disallowed.
|
| 51 |
+
* 2026-04-15:
|
| 52 |
+
* Initial version.
|
| 53 |
+
|
| 54 |
+
Contributors from the Kaggle Community:
|
| 55 |
+
* @anglolodorf
|
| 56 |
+
* @arc144
|
| 57 |
+
* @asalhi
|
| 58 |
+
* @calibrator
|
| 59 |
+
* @cdeotte
|
| 60 |
+
* @hengck23
|
| 61 |
+
* @jazivxt
|
| 62 |
+
* @jiweiliu
|
| 63 |
+
* @kameronkilchrist
|
| 64 |
+
* @kevinyuluo
|
| 65 |
+
* @kosirowada
|
| 66 |
+
* @linkinpony
|
| 67 |
+
* @maxjeblick
|
| 68 |
+
* @mukundan314
|
| 69 |
+
* @pavelsavchenkov
|
| 70 |
+
* @prokaj
|
| 71 |
+
* @robga
|
| 72 |
+
* @shinh0
|
| 73 |
+
* @tonylica
|
| 74 |
+
* @yeoyunsianggeremie
|
| 75 |
+
* @yiheng
|
| 76 |
+
"""
|
| 77 |
+
|
| 78 |
+
import itertools
|
| 79 |
+
import json
|
| 80 |
+
import math
|
| 81 |
+
import pathlib
|
| 82 |
+
import traceback
|
| 83 |
+
|
| 84 |
+
import IPython.display
|
| 85 |
+
import matplotlib.pyplot as plt
|
| 86 |
+
import numpy as np
|
| 87 |
+
import onnx
|
| 88 |
+
import onnx_tool
|
| 89 |
+
import onnxruntime
|
| 90 |
+
|
| 91 |
+
|
| 92 |
+
display = IPython.display.display
|
| 93 |
+
FileLink = IPython.display.FileLink
|
| 94 |
+
|
| 95 |
+
_BATCH_SIZE, _CHANNELS, _HEIGHT, _WIDTH = 1, 10, 30, 30
|
| 96 |
+
_NEUROGOLF_DIR = "/kaggle/input/competitions/neurogolf-2026/"
|
| 97 |
+
_COLORS = [
|
| 98 |
+
(0, 0, 0),
|
| 99 |
+
(30, 147, 255),
|
| 100 |
+
(250, 61, 49),
|
| 101 |
+
(78, 204, 48),
|
| 102 |
+
(255, 221, 0),
|
| 103 |
+
(153, 153, 153),
|
| 104 |
+
(229, 59, 163),
|
| 105 |
+
(255, 133, 28),
|
| 106 |
+
(136, 216, 241),
|
| 107 |
+
(147, 17, 49),
|
| 108 |
+
(240, 240, 240),
|
| 109 |
+
(146, 117, 86)
|
| 110 |
+
]
|
| 111 |
+
_DATA_TYPE = onnx.TensorProto.FLOAT
|
| 112 |
+
_EXCLUDED_OP_TYPES = ["LOOP", "SCAN", "NONZERO", "UNIQUE", "SCRIPT", "FUNCTION", "COMPRESS"]
|
| 113 |
+
_FILESIZE_LIMIT_IN_BYTES = 1.44 * 1024 * 1024
|
| 114 |
+
_GRID_SHAPE = [_BATCH_SIZE, _CHANNELS, _HEIGHT, _WIDTH]
|
| 115 |
+
_IR_VERSION, _OPSET_IMPORTS = 10, [onnx.helper.make_opsetid("", 10)]
|
| 116 |
+
_TASK_ZERO = {
|
| 117 |
+
"train": [{
|
| 118 |
+
"input": [
|
| 119 |
+
[5, 5, 5, 5, 5, 5, 5, 5, 5, 5],
|
| 120 |
+
[5, 1, 1, 1, 1, 1, 1, 5, 5, 5],
|
| 121 |
+
[5, 1, 1, 1, 1, 1, 1, 5, 5, 5],
|
| 122 |
+
[5, 1, 1, 1, 1, 1, 1, 5, 5, 5],
|
| 123 |
+
[5, 1, 1, 1, 1, 1, 1, 5, 5, 5],
|
| 124 |
+
[5, 1, 1, 1, 1, 1, 1, 5, 5, 5],
|
| 125 |
+
[5, 1, 1, 1, 1, 1, 1, 5, 5, 5],
|
| 126 |
+
[5, 5, 5, 5, 5, 5, 5, 5, 5, 5],
|
| 127 |
+
[5, 5, 5, 5, 5, 5, 5, 5, 5, 5],
|
| 128 |
+
[5, 5, 5, 5, 5, 5, 5, 5, 5, 5],
|
| 129 |
+
],
|
| 130 |
+
"output": [
|
| 131 |
+
[5, 5, 5, 5, 5, 5, 5, 5, 5, 5],
|
| 132 |
+
[5, 1, 1, 1, 1, 1, 1, 5, 5, 5],
|
| 133 |
+
[5, 1, 1, 1, 1, 1, 1, 0, 5, 5],
|
| 134 |
+
[5, 1, 1, 1, 1, 1, 1, 0, 5, 5],
|
| 135 |
+
[5, 1, 1, 1, 1, 1, 1, 0, 5, 5],
|
| 136 |
+
[5, 1, 1, 1, 1, 1, 1, 0, 5, 5],
|
| 137 |
+
[5, 1, 1, 1, 1, 1, 1, 0, 5, 5],
|
| 138 |
+
[5, 5, 0, 0, 0, 0, 0, 0, 5, 5],
|
| 139 |
+
[5, 5, 5, 5, 5, 5, 5, 5, 5, 5],
|
| 140 |
+
[5, 5, 5, 5, 5, 5, 5, 5, 5, 5],
|
| 141 |
+
],
|
| 142 |
+
}],
|
| 143 |
+
"test": [{
|
| 144 |
+
"input": [
|
| 145 |
+
[5, 5, 5, 5, 5, 5, 5, 5, 5, 5],
|
| 146 |
+
[5, 5, 4, 4, 4, 4, 4, 4, 5, 5],
|
| 147 |
+
[5, 5, 4, 4, 4, 4, 4, 4, 5, 5],
|
| 148 |
+
[5, 5, 5, 5, 5, 5, 5, 5, 5, 5],
|
| 149 |
+
[5, 5, 5, 5, 5, 5, 5, 5, 5, 5],
|
| 150 |
+
[5, 5, 4, 4, 4, 4, 4, 5, 5, 5],
|
| 151 |
+
[5, 5, 4, 5, 5, 5, 4, 5, 5, 5],
|
| 152 |
+
[5, 5, 4, 5, 5, 5, 4, 5, 5, 5],
|
| 153 |
+
[5, 5, 4, 4, 4, 4, 4, 5, 5, 5],
|
| 154 |
+
[5, 5, 5, 5, 5, 5, 5, 5, 5, 5],
|
| 155 |
+
],
|
| 156 |
+
"output": [
|
| 157 |
+
[5, 5, 5, 5, 5, 5, 5, 5, 5, 5],
|
| 158 |
+
[5, 5, 4, 4, 4, 4, 4, 4, 5, 5],
|
| 159 |
+
[5, 5, 4, 4, 4, 4, 4, 4, 0, 5],
|
| 160 |
+
[5, 5, 5, 0, 0, 0, 0, 0, 0, 5],
|
| 161 |
+
[5, 5, 5, 5, 5, 5, 5, 5, 5, 5],
|
| 162 |
+
[5, 5, 4, 4, 4, 4, 4, 5, 5, 5],
|
| 163 |
+
[5, 5, 4, 0, 0, 0, 4, 0, 5, 5],
|
| 164 |
+
[5, 5, 4, 0, 5, 5, 4, 0, 5, 5],
|
| 165 |
+
[5, 5, 4, 4, 4, 4, 4, 0, 5, 5],
|
| 166 |
+
[5, 5, 5, 0, 0, 0, 0, 0, 5, 5],
|
| 167 |
+
],
|
| 168 |
+
}],
|
| 169 |
+
"arc-gen": [{
|
| 170 |
+
"input": [
|
| 171 |
+
[5, 5, 5, 5, 5, 5, 5, 5, 5, 5],
|
| 172 |
+
[5, 5, 5, 5, 5, 5, 5, 5, 5, 5],
|
| 173 |
+
[5, 5, 2, 2, 2, 2, 2, 2, 5, 5],
|
| 174 |
+
[5, 5, 2, 5, 5, 5, 5, 2, 5, 5],
|
| 175 |
+
[5, 5, 2, 5, 5, 5, 5, 2, 5, 5],
|
| 176 |
+
[5, 5, 2, 5, 5, 5, 5, 2, 5, 5],
|
| 177 |
+
[5, 5, 2, 5, 5, 5, 5, 2, 5, 5],
|
| 178 |
+
[5, 5, 2, 2, 2, 2, 2, 2, 5, 5],
|
| 179 |
+
[5, 5, 5, 5, 5, 5, 5, 5, 5, 5],
|
| 180 |
+
[5, 5, 5, 5, 5, 5, 5, 5, 5, 5],
|
| 181 |
+
],
|
| 182 |
+
"output": [
|
| 183 |
+
[5, 5, 5, 5, 5, 5, 5, 5, 5, 5],
|
| 184 |
+
[5, 5, 5, 5, 5, 5, 5, 5, 5, 5],
|
| 185 |
+
[5, 5, 2, 2, 2, 2, 2, 2, 5, 5],
|
| 186 |
+
[5, 5, 2, 0, 0, 0, 0, 2, 0, 5],
|
| 187 |
+
[5, 5, 2, 0, 5, 5, 5, 2, 0, 5],
|
| 188 |
+
[5, 5, 2, 0, 5, 5, 5, 2, 0, 5],
|
| 189 |
+
[5, 5, 2, 0, 5, 5, 5, 2, 0, 5],
|
| 190 |
+
[5, 5, 2, 2, 2, 2, 2, 2, 0, 5],
|
| 191 |
+
[5, 5, 5, 0, 0, 0, 0, 0, 0, 5],
|
| 192 |
+
[5, 5, 5, 5, 5, 5, 5, 5, 5, 5],
|
| 193 |
+
],
|
| 194 |
+
}],
|
| 195 |
+
}
|
| 196 |
+
|
| 197 |
+
|
| 198 |
+
def calculate_memory(model, trace_path):
|
| 199 |
+
onnx.checker.check_model(model, full_check=True)
|
| 200 |
+
graph = onnx.shape_inference.infer_shapes(model, strict_mode=True).graph
|
| 201 |
+
if len(graph.input) > 1 or len(graph.output) > 1: return None
|
| 202 |
+
init_names = {init.name for init in graph.initializer}
|
| 203 |
+
init_names.update(init.name for init in graph.sparse_initializer)
|
| 204 |
+
io_names = {t.name for t in list(graph.input) + list(graph.output)}
|
| 205 |
+
if io_names.intersection(init_names): return None
|
| 206 |
+
if model.functions: return None
|
| 207 |
+
for opset in model.opset_import:
|
| 208 |
+
if opset.domain not in {"", "ai.onnx"}: return None
|
| 209 |
+
node_outputs = {}
|
| 210 |
+
tensor_names = set()
|
| 211 |
+
for node in graph.node:
|
| 212 |
+
for attr in node.attribute:
|
| 213 |
+
if attr.type in [onnx.AttributeProto.GRAPH,
|
| 214 |
+
onnx.AttributeProto.GRAPHS]:
|
| 215 |
+
return None
|
| 216 |
+
node_outputs[node.name] = list(node.output)
|
| 217 |
+
for output_name in node.output:
|
| 218 |
+
if output_name: tensor_names.add(output_name)
|
| 219 |
+
tensor_memory = {}
|
| 220 |
+
tensor_dtypes = {}
|
| 221 |
+
tensor_map = {
|
| 222 |
+
t.name: t for t in list(graph.input) + list(graph.value_info) + list(graph.output)
|
| 223 |
+
}
|
| 224 |
+
tensor_names.update(tensor_map.keys())
|
| 225 |
+
for tensor_name in tensor_names:
|
| 226 |
+
item = tensor_map.get(tensor_name)
|
| 227 |
+
if not item: return None
|
| 228 |
+
if item.type.HasField("sequence_type"): return None
|
| 229 |
+
if not item.type.HasField("tensor_type"): continue
|
| 230 |
+
tensor_type = item.type.tensor_type
|
| 231 |
+
if not tensor_type.HasField("shape"): return None
|
| 232 |
+
num_elements = 1
|
| 233 |
+
for dim in tensor_type.shape.dim:
|
| 234 |
+
if dim.HasField("dim_param"): return None
|
| 235 |
+
if not dim.HasField("dim_value"): return None
|
| 236 |
+
if dim.dim_value <= 0: return None
|
| 237 |
+
num_elements *= dim.dim_value
|
| 238 |
+
if tensor_name in ['input', 'output']: continue
|
| 239 |
+
np_dtype = onnx.helper.tensor_dtype_to_np_dtype(tensor_type.elem_type)
|
| 240 |
+
tensor_memory[tensor_name] = num_elements * np.dtype(np_dtype).itemsize
|
| 241 |
+
tensor_dtypes[tensor_name] = np_dtype
|
| 242 |
+
|
| 243 |
+
# Defensive check to verify uniqueness.
|
| 244 |
+
seen = set()
|
| 245 |
+
for item in list(graph.input) + list(graph.value_info) + list(graph.output):
|
| 246 |
+
if item.name in seen: return None
|
| 247 |
+
seen.add(item.name)
|
| 248 |
+
for node in graph.node:
|
| 249 |
+
for output_name in node.output:
|
| 250 |
+
if output_name and output_name != "output":
|
| 251 |
+
item = tensor_map.get(output_name)
|
| 252 |
+
if item is None or not item.type.HasField("tensor_type"):
|
| 253 |
+
return None
|
| 254 |
+
|
| 255 |
+
# Retrieve actual tensor shapes via the ONNX Runtime Profiler's JSON Trace.
|
| 256 |
+
with open(trace_path, 'r') as f:
|
| 257 |
+
trace_data = json.load(f)
|
| 258 |
+
for event in trace_data:
|
| 259 |
+
if event.get("cat") != "Node" or "args" not in event: continue
|
| 260 |
+
if "output_type_shape" not in event["args"]: continue
|
| 261 |
+
node_name = event.get("name").replace("_kernel_time", "")
|
| 262 |
+
if node_name not in node_outputs: continue
|
| 263 |
+
for i, shape_dict in enumerate(event["args"]["output_type_shape"]):
|
| 264 |
+
if i >= len(node_outputs[node_name]): continue
|
| 265 |
+
output_name = node_outputs[node_name][i]
|
| 266 |
+
if output_name not in tensor_dtypes: continue
|
| 267 |
+
itemsize = np.dtype(tensor_dtypes[output_name]).itemsize
|
| 268 |
+
mem = itemsize * sum(math.prod(dims) for dims in shape_dict.values())
|
| 269 |
+
tensor_memory[output_name] = max(tensor_memory[output_name], mem)
|
| 270 |
+
return sum(tensor_memory.values())
|
| 271 |
+
|
| 272 |
+
|
| 273 |
+
def check_network(filename):
|
| 274 |
+
file_path = pathlib.Path(filename)
|
| 275 |
+
if not file_path.is_file():
|
| 276 |
+
print(f"Error: File {filename} does not exist.")
|
| 277 |
+
return False
|
| 278 |
+
if (filesize := file_path.stat().st_size) > _FILESIZE_LIMIT_IN_BYTES:
|
| 279 |
+
print(f"Error: Filesize {filesize} exceeds {_FILESIZE_LIMIT_IN_BYTES}.")
|
| 280 |
+
return False
|
| 281 |
+
return True
|
| 282 |
+
|
| 283 |
+
|
| 284 |
+
def convert_to_numpy(example):
|
| 285 |
+
benchmark = {}
|
| 286 |
+
example_shape = (1, _CHANNELS, _HEIGHT, _WIDTH)
|
| 287 |
+
for mode in ["input", "output"]:
|
| 288 |
+
benchmark[mode] = np.zeros(example_shape, dtype=np.float32)
|
| 289 |
+
grid = example[mode]
|
| 290 |
+
if max(len(grid), len(grid[0])) > 30: return None
|
| 291 |
+
for r, _ in enumerate(grid):
|
| 292 |
+
for c, color in enumerate(grid[r]):
|
| 293 |
+
benchmark[mode][0][color][r][c] = 1.0
|
| 294 |
+
return benchmark
|
| 295 |
+
|
| 296 |
+
|
| 297 |
+
def convert_from_numpy(benchmark):
|
| 298 |
+
example = []
|
| 299 |
+
_, channels, height, width = benchmark.shape
|
| 300 |
+
for row in range(height):
|
| 301 |
+
cells = []
|
| 302 |
+
for col in range(width):
|
| 303 |
+
colors = [c for c in range(channels) if benchmark[0][c][row][col] == 1]
|
| 304 |
+
cells.append(colors[0] if len(colors) == 1 else (11 if colors else 10))
|
| 305 |
+
while cells and cells[-1] == 10:
|
| 306 |
+
cells.pop(-1)
|
| 307 |
+
example.append(cells)
|
| 308 |
+
while example and not example[-1]:
|
| 309 |
+
example.pop(-1)
|
| 310 |
+
return example
|
| 311 |
+
|
| 312 |
+
|
| 313 |
+
def calculate_params(model):
|
| 314 |
+
params = 0
|
| 315 |
+
for init in model.graph.initializer:
|
| 316 |
+
if any(d <= 0 for d in init.dims): return None
|
| 317 |
+
params += math.prod(init.dims)
|
| 318 |
+
for sparse_init in model.graph.sparse_initializer:
|
| 319 |
+
if any(d <= 0 for d in sparse_init.values.dims): return None
|
| 320 |
+
params += math.prod(sparse_init.values.dims)
|
| 321 |
+
for node in model.graph.node:
|
| 322 |
+
if node.op_type != 'Constant': continue
|
| 323 |
+
for attr in node.attribute:
|
| 324 |
+
if attr.name == 'value':
|
| 325 |
+
if any(d <= 0 for d in attr.t.dims): return None
|
| 326 |
+
params += math.prod(attr.t.dims)
|
| 327 |
+
elif attr.name == 'sparse_value':
|
| 328 |
+
if any(d <= 0 for d in attr.sparse_tensor.values.dims): return None
|
| 329 |
+
params += math.prod(attr.sparse_tensor.values.dims)
|
| 330 |
+
elif attr.name == 'value_floats':
|
| 331 |
+
params += len(attr.floats)
|
| 332 |
+
elif attr.name == 'value_ints':
|
| 333 |
+
params += len(attr.ints)
|
| 334 |
+
elif attr.name == 'value_strings':
|
| 335 |
+
params += len(attr.strings)
|
| 336 |
+
return params
|
| 337 |
+
|
| 338 |
+
|
| 339 |
+
def score_network(sanitized, trace_path):
|
| 340 |
+
for node in sanitized.graph.node:
|
| 341 |
+
if node.op_type.upper() in _EXCLUDED_OP_TYPES:
|
| 342 |
+
print(f"Error: Op type {node.op_type} is not permitted.")
|
| 343 |
+
return None, None
|
| 344 |
+
if "Sequence" in node.op_type:
|
| 345 |
+
print(f"Error: Op type {node.op_type} is not permitted.")
|
| 346 |
+
return None, None
|
| 347 |
+
return calculate_memory(sanitized, trace_path), calculate_params(sanitized)
|
| 348 |
+
|
| 349 |
+
|
| 350 |
+
def sanitize_model(model):
|
| 351 |
+
for node in model.graph.node:
|
| 352 |
+
node.name = node.output[0]
|
| 353 |
+
if "kernel_time" in node.output[0]: return None
|
| 354 |
+
|
| 355 |
+
name_map, counter = {}, 0
|
| 356 |
+
|
| 357 |
+
def get_safe_name(old_name):
|
| 358 |
+
nonlocal counter
|
| 359 |
+
if not old_name or old_name in ["input", "output"]: return old_name
|
| 360 |
+
if old_name not in name_map:
|
| 361 |
+
name_map[old_name] = f"safe_name_{counter}"
|
| 362 |
+
counter += 1
|
| 363 |
+
return name_map[old_name]
|
| 364 |
+
|
| 365 |
+
for inp in model.graph.input:
|
| 366 |
+
inp.name = get_safe_name(inp.name)
|
| 367 |
+
for init in model.graph.initializer:
|
| 368 |
+
init.name = get_safe_name(init.name)
|
| 369 |
+
|
| 370 |
+
for node in model.graph.node:
|
| 371 |
+
for i in range(len(node.input)):
|
| 372 |
+
node.input[i] = get_safe_name(node.input[i])
|
| 373 |
+
for i in range(len(node.output)):
|
| 374 |
+
node.output[i] = get_safe_name(node.output[i])
|
| 375 |
+
if len(node.output) > 0 and node.output[0]:
|
| 376 |
+
node.name = node.output[0]
|
| 377 |
+
|
| 378 |
+
for out in model.graph.output:
|
| 379 |
+
out.name = get_safe_name(out.name)
|
| 380 |
+
for vi in model.graph.value_info:
|
| 381 |
+
vi.name = get_safe_name(vi.name)
|
| 382 |
+
for node in model.graph.node:
|
| 383 |
+
node.name = node.output[0]
|
| 384 |
+
return model
|
| 385 |
+
|
| 386 |
+
|
| 387 |
+
def load_examples(task_num):
|
| 388 |
+
"""Loads relevant data from ARC-AGI and ARC-GEN."""
|
| 389 |
+
if not task_num:
|
| 390 |
+
return _TASK_ZERO
|
| 391 |
+
with open(_NEUROGOLF_DIR + f"task{task_num:03d}.json") as f:
|
| 392 |
+
examples = json.load(f)
|
| 393 |
+
return examples
|
| 394 |
+
|
| 395 |
+
|
| 396 |
+
def run_network(session, benchmark_input):
|
| 397 |
+
result = session.run(["output"], {"input": benchmark_input})
|
| 398 |
+
return (result[0] > 0.0).astype(float)
|
| 399 |
+
|
| 400 |
+
|
| 401 |
+
def show_examples(examples, bgcolor=(255, 255, 255)):
|
| 402 |
+
# Determine the dimensions of the image to be rendered.
|
| 403 |
+
width, height, offset = 0, 0, 1
|
| 404 |
+
for example in examples:
|
| 405 |
+
grid, output = example["input"], example["output"]
|
| 406 |
+
width += len(grid[0]) + 1 + len(output[0]) + 4
|
| 407 |
+
height = max(height, max(len(grid), len(output)) + 4)
|
| 408 |
+
# Determine the contents of the image.
|
| 409 |
+
image = [[bgcolor for _ in range(width)] for _ in range(height)]
|
| 410 |
+
for example in examples:
|
| 411 |
+
grid, output = example["input"], example["output"]
|
| 412 |
+
grid_width, output_width = len(grid[0]), len(output[0])
|
| 413 |
+
for r, row in enumerate(grid):
|
| 414 |
+
for c, cell in enumerate(row):
|
| 415 |
+
image[r + 2][offset + c + 1] = _COLORS[cell]
|
| 416 |
+
offset += grid_width + 1
|
| 417 |
+
for r, row in enumerate(output):
|
| 418 |
+
for c, cell in enumerate(row):
|
| 419 |
+
image[r + 2][offset + c + 1] = _COLORS[cell]
|
| 420 |
+
offset += output_width + 4
|
| 421 |
+
# Draw the image.
|
| 422 |
+
fig = plt.figure(figsize=(10, 5))
|
| 423 |
+
ax = fig.add_axes([0, 0, 1, 1])
|
| 424 |
+
ax.imshow(np.array(image))
|
| 425 |
+
# Draw the horizontal and vertical lines.
|
| 426 |
+
offset = 1
|
| 427 |
+
for example in examples:
|
| 428 |
+
grid, output = example["input"], example["output"]
|
| 429 |
+
grid_width, grid_height = len(grid[0]), len(grid)
|
| 430 |
+
output_width, output_height = len(output[0]), len(output)
|
| 431 |
+
ax.hlines([r + 1.5 for r in range(grid_height+1)],
|
| 432 |
+
xmin=offset+0.5, xmax=offset+grid_width+0.5, color="black")
|
| 433 |
+
ax.vlines([offset + c + 0.5 for c in range(grid_width+1)],
|
| 434 |
+
ymin=1.5, ymax=grid_height+1.5, color="black")
|
| 435 |
+
offset += grid_width + 1
|
| 436 |
+
ax.hlines([r + 1.5 for r in range(output_height+1)],
|
| 437 |
+
xmin=offset+0.5, xmax=offset+output_width+0.5, color="black")
|
| 438 |
+
ax.vlines([offset + c + 0.5 for c in range(output_width+1)],
|
| 439 |
+
ymin=1.5, ymax=output_height+1.5, color="black")
|
| 440 |
+
offset += output_width + 2
|
| 441 |
+
ax.vlines([offset+0.5], ymin=-0.5, ymax=height-0.5, color="black")
|
| 442 |
+
offset += 2
|
| 443 |
+
ax.set_xticks([])
|
| 444 |
+
ax.set_yticks([])
|
| 445 |
+
|
| 446 |
+
|
| 447 |
+
def show_legend():
|
| 448 |
+
image = [[(255, 255, 255) for _ in range(21)] for _ in range(5)]
|
| 449 |
+
for idx, color in enumerate(_COLORS[:10]):
|
| 450 |
+
image[1][2 * idx + 1] = color
|
| 451 |
+
for idx, color in enumerate(_COLORS[10:]):
|
| 452 |
+
for col in range(3):
|
| 453 |
+
image[3][12 * idx + col + 3] = color
|
| 454 |
+
fig = plt.figure(figsize=(10, 5))
|
| 455 |
+
ax = fig.add_axes([0, 0, 1, 1])
|
| 456 |
+
ax.imshow(np.array(image))
|
| 457 |
+
for idx, _ in enumerate(_COLORS[:10]):
|
| 458 |
+
color = "white" if idx in [0, 9] else "black"
|
| 459 |
+
ax.text(2 * idx + 0.9, 1.1, str(idx), color=color)
|
| 460 |
+
ax.text(3.4, 3.1, "no color", color="black")
|
| 461 |
+
ax.text(5.75, 3.1, "<--- special colors to indicate one-hot encoding errors --->", color="black")
|
| 462 |
+
ax.text(14.85, 3.1, "too many colors", color="white")
|
| 463 |
+
ax.set_xticks([])
|
| 464 |
+
ax.set_yticks([])
|
| 465 |
+
|
| 466 |
+
|
| 467 |
+
def single_layer_conv2d_network(weight_fn, kernel_size):
|
| 468 |
+
kernel_offsets = range(-kernel_size // 2 + 1, kernel_size // 2 + 1)
|
| 469 |
+
kernel_shape = [kernel_size, kernel_size]
|
| 470 |
+
w_shape = [_CHANNELS, _CHANNELS, kernel_size, kernel_size]
|
| 471 |
+
pads = [kernel_size // 2] * 4
|
| 472 |
+
weight_cells = itertools.product(range(_CHANNELS), range(_CHANNELS),
|
| 473 |
+
kernel_offsets, kernel_offsets)
|
| 474 |
+
weights = [weight_fn(o, i, (r, c)) for (o, i, r, c) in weight_cells]
|
| 475 |
+
|
| 476 |
+
x = onnx.helper.make_tensor_value_info("input", _DATA_TYPE, _GRID_SHAPE)
|
| 477 |
+
y = onnx.helper.make_tensor_value_info("output", _DATA_TYPE, _GRID_SHAPE)
|
| 478 |
+
w = onnx.helper.make_tensor("W", _DATA_TYPE, w_shape, weights)
|
| 479 |
+
node_def = onnx.helper.make_node("Conv", ["input", "W"], ["output"],
|
| 480 |
+
kernel_shape=kernel_shape, pads=pads)
|
| 481 |
+
graph_def = onnx.helper.make_graph([node_def], "graph", [x], [y], [w])
|
| 482 |
+
model_def = onnx.helper.make_model(graph_def, ir_version=_IR_VERSION,
|
| 483 |
+
opset_imports=_OPSET_IMPORTS)
|
| 484 |
+
return model_def
|
| 485 |
+
|
| 486 |
+
|
| 487 |
+
def verify_network(network, task_num, examples):
|
| 488 |
+
filename = "task{:03d}.onnx".format(task_num)
|
| 489 |
+
onnx.save(network, filename)
|
| 490 |
+
if not check_network(filename): return
|
| 491 |
+
try:
|
| 492 |
+
# Load the model, sanitize node names, and enable profiling.
|
| 493 |
+
sanitized = sanitize_model(onnx.load(filename))
|
| 494 |
+
if not sanitized: return
|
| 495 |
+
options = onnxruntime.SessionOptions()
|
| 496 |
+
options.enable_profiling = True
|
| 497 |
+
options.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_DISABLE_ALL
|
| 498 |
+
options.profile_file_prefix = f"{task_num:03}"
|
| 499 |
+
session = onnxruntime.InferenceSession(sanitized.SerializeToString(), options)
|
| 500 |
+
except onnxruntime.ONNXRuntimeError as e:
|
| 501 |
+
print(f"Error: Unable to load ONNX model: {e}")
|
| 502 |
+
return
|
| 503 |
+
arc_agi_right, arc_agi_wrong, arc_agi_expected = verify_subset(session, examples["train"] + examples["test"])
|
| 504 |
+
arc_gen_right, arc_gen_wrong, arc_gen_expected = verify_subset(session, examples["arc-gen"])
|
| 505 |
+
print(f"Results on ARC-AGI examples: {arc_agi_right} pass, {arc_agi_wrong} fail")
|
| 506 |
+
print(f"Results on ARC-GEN examples: {arc_gen_right} pass, {arc_gen_wrong} fail")
|
| 507 |
+
print()
|
| 508 |
+
memory, params = score_network(sanitized, session.end_profiling())
|
| 509 |
+
if memory is None or params is None:
|
| 510 |
+
print("Error: Your network performance could not be measured")
|
| 511 |
+
if memory < 0 or params < 0:
|
| 512 |
+
print("Error: Your network performance could not be measured")
|
| 513 |
+
elif arc_agi_wrong + arc_gen_wrong == 0:
|
| 514 |
+
print("Your network IS READY for submission!")
|
| 515 |
+
print()
|
| 516 |
+
print("Performance stats (memory values reported here are approximate):")
|
| 517 |
+
onnx_tool.model_profile(filename)
|
| 518 |
+
points = max(1.0, 25.0 - math.log(max(1.0, memory + params)))
|
| 519 |
+
print()
|
| 520 |
+
print(f"It appears to require {memory} bytes + {params} params, yielding {points:.3f} points.")
|
| 521 |
+
print()
|
| 522 |
+
print("Next steps:")
|
| 523 |
+
print(f" * Click the link below to download {filename} onto your local machine.")
|
| 524 |
+
print(" * Create a zip file containing that network along with all others.")
|
| 525 |
+
print(" * Submit that zip file to the Kaggle competition so that it can be officially scored.")
|
| 526 |
+
print()
|
| 527 |
+
display(FileLink(filename))
|
| 528 |
+
else:
|
| 529 |
+
print("Your network IS NOT ready for submission.")
|
| 530 |
+
expected = None
|
| 531 |
+
expected = arc_agi_expected if arc_agi_expected is not None else expected
|
| 532 |
+
expected = arc_gen_expected if arc_gen_expected is not None else expected
|
| 533 |
+
if expected is None: return
|
| 534 |
+
benchmark = convert_to_numpy(expected)
|
| 535 |
+
actual = {}
|
| 536 |
+
actual["input"] = expected["input"]
|
| 537 |
+
actual["output"] = convert_from_numpy(run_network(session, benchmark["input"]))
|
| 538 |
+
print("The expected result is shown in green; your actual result is shown in red.")
|
| 539 |
+
show_examples([expected], bgcolor=(200, 255, 200))
|
| 540 |
+
show_examples([actual], bgcolor=(255, 200, 200))
|
| 541 |
+
|
| 542 |
+
|
| 543 |
+
def verify_subset(session, example_subset):
|
| 544 |
+
right, wrong, expected, error = 0, 0, None, ""
|
| 545 |
+
for example in example_subset:
|
| 546 |
+
benchmark = convert_to_numpy(example)
|
| 547 |
+
if not benchmark: continue
|
| 548 |
+
try:
|
| 549 |
+
user_output = run_network(session, benchmark["input"])
|
| 550 |
+
if np.array_equal(user_output, benchmark["output"]):
|
| 551 |
+
right += 1
|
| 552 |
+
else:
|
| 553 |
+
expected = example
|
| 554 |
+
wrong += 1
|
| 555 |
+
except onnxruntime.ONNXRuntimeError:
|
| 556 |
+
error = traceback.format_exc()
|
| 557 |
+
wrong += 1
|
| 558 |
+
if error: print(f"Error: {error}")
|
| 559 |
+
return right, wrong, expected
|