File size: 5,298 Bytes
81796cb | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 | #!/usr/bin/env python3
"""
fix_scoring_bugs.py — Fix ONNX models with scoring bugs.
Bug 1: dims=[0] initializers (scalar_shape)
- The scorer rejects any initializer with dims containing 0
- Fix: Replace Reshape(x, scalar_shape) with Squeeze(x)
Bug 2: Missing value_info (only fixable if shapes are static)
- Run shape inference to populate value_info
- Only works if model has no dynamic-shape ops (NonZero, Where with data-dependent shapes)
Usage:
python fix_scoring_bugs.py --input task291.onnx --output fixed/task291.onnx --task-num 291
"""
import argparse
import os
import sys
import onnx
import onnxruntime as ort
import numpy as np
import json
from onnx import helper, TensorProto, numpy_helper, shape_inference
def fix_scalar_shape(model):
"""Replace Reshape(x, scalar_shape) with Squeeze(x) where scalar_shape has dims=[0]."""
# Find problematic initializers
bad_inits = {init.name for init in model.graph.initializer if any(d <= 0 for d in init.dims)}
if not bad_inits:
return model, False
# Remove bad initializers
new_inits = [init for init in model.graph.initializer if init.name not in bad_inits]
# Replace Reshape nodes that use bad initializers with Squeeze
new_nodes = []
fixed = False
for node in model.graph.node:
if node.op_type == 'Reshape' and any(inp in bad_inits for inp in node.input):
squeeze_node = helper.make_node(
'Squeeze',
inputs=[node.input[0]],
outputs=list(node.output),
)
new_nodes.append(squeeze_node)
fixed = True
else:
new_nodes.append(node)
del model.graph.initializer[:]
model.graph.initializer.extend(new_inits)
del model.graph.node[:]
model.graph.node.extend(new_nodes)
return model, fixed
def fix_missing_value_info(model):
"""Run shape inference to add value_info for models that lack it."""
if len(model.graph.value_info) > 0:
return model, False
try:
inferred = shape_inference.infer_shapes(model, strict_mode=True)
if len(inferred.graph.value_info) > 0:
return inferred, True
except:
pass
try:
inferred = shape_inference.infer_shapes(model, strict_mode=False)
if len(inferred.graph.value_info) > 0:
return inferred, True
except:
pass
return model, False
def validate_model(model_path, task_num, task_data_dir):
"""Validate model produces correct outputs for all examples."""
task_json = os.path.join(task_data_dir, f'task{task_num:03d}.json')
if not os.path.exists(task_json):
print(f" Warning: task data not found at {task_json}")
return True # Can't validate without data
sess = ort.InferenceSession(model_path)
with open(task_json) as f:
data = json.load(f)
right, wrong = 0, 0
for ex in data['train'] + data['test'] + data.get('arc-gen', []):
inp = np.zeros((1, 10, 30, 30), dtype=np.float32)
for r, row in enumerate(ex['input']):
for c, v in enumerate(row):
if r < 30 and c < 30:
inp[0][v][r][c] = 1.0
result = sess.run(['output'], {'input': inp})
out = (result[0] > 0.0).astype(float)
exp = np.zeros((1, 10, 30, 30), dtype=np.float32)
for r, row in enumerate(ex['output']):
for c, v in enumerate(row):
if r < 30 and c < 30:
exp[0][v][r][c] = 1.0
if np.array_equal(out, exp):
right += 1
else:
wrong += 1
print(f" Validation: {right} pass, {wrong} fail")
return wrong == 0
def main():
parser = argparse.ArgumentParser(description="Fix ONNX scoring bugs")
parser.add_argument('--input', required=True, help='Input ONNX model')
parser.add_argument('--output', required=True, help='Output ONNX model')
parser.add_argument('--task-num', type=int, required=True, help='Task number')
parser.add_argument('--task-data-dir', default='/app/task-data', help='Task data directory')
args = parser.parse_args()
print(f"Loading {args.input}...")
model = onnx.load(args.input)
print(f" Nodes: {len(model.graph.node)}")
print(f" Initializers: {len(model.graph.initializer)}")
print(f" Value_info: {len(model.graph.value_info)}")
# Apply fixes
model, fixed_scalar = fix_scalar_shape(model)
if fixed_scalar:
print(" ✓ Fixed scalar_shape (dims=[0]) bug")
model, fixed_vi = fix_missing_value_info(model)
if fixed_vi:
print(f" ✓ Added value_info via shape inference ({len(model.graph.value_info)} entries)")
if not fixed_scalar and not fixed_vi:
print(" No fixes needed")
return
# Save
os.makedirs(os.path.dirname(args.output) or '.', exist_ok=True)
onnx.save(model, args.output)
print(f" Saved to {args.output}")
# Validate
if validate_model(args.output, args.task_num, args.task_data_dir):
print(" ✓ Model passes all examples!")
else:
print(" ✗ Model FAILS validation — do NOT use!")
sys.exit(1)
if __name__ == '__main__':
main()
|