rogermt commited on
Commit
f907ee5
·
verified ·
1 Parent(s): e323bad

Upload neurogolf_utils.py

Browse files
Files changed (1) hide show
  1. 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