File size: 4,906 Bytes
bdf0499 | 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 | """Build optimized ONNX model for Task 084.
Algorithm (175/175 verified):
1. Input has a colored column at col 0 (N rows tall = grid is NxN)
2. Draw anti-diagonal (color 2) from (0, N-1) to (N-2, 1)
3. Fill bottom row (N-1, cols 1..N-1) with color 4
4. Keep original column 0
ONNX approach:
- Detect N from sum of non-bg channels
- Use row/col grids to compute masks for anti-diagonal and bottom row
- Combine into output using channel selection
Score: 175/175 pass, file size 10KB (~12 pts estimated)
Usage:
python build_task084_onnx.py --task-data-dir ../task-data
"""
from onnx import TensorProto
import numpy as np
from onnx_builder import OnnxBuilder, build_and_validate
def build_task084():
b = OnnxBuilder()
const, nd = b.const, b.nd
const('c_one', [1.0])
const('c_zero', [0.0])
const('c_half', [0.5])
const('axes123', [1, 2, 3], 'i')
const('axes23', [2, 3], 'i')
const('axes1', [1], 'i')
const('shape_1_1_30_30', [1, 1, 30, 30], 'i')
const('shape_1_1_1_1', [1, 1, 1, 1], 'i')
const('shape_1_10_1_1', [1, 10, 1, 1], 'i')
row_vals = np.arange(30, dtype=np.float32).reshape(1, 1, 30, 1) * np.ones((1, 1, 1, 30), dtype=np.float32)
col_vals = np.arange(30, dtype=np.float32).reshape(1, 1, 1, 30) * np.ones((1, 1, 30, 1), dtype=np.float32)
const('row_grid', row_vals)
const('col_grid', col_vals)
const('idx_0', [0], 'i')
const('idx_1', [1], 'i')
const('idx_2', [2], 'i')
const('idx_10', [10], 'i')
const('depth_10', [10.0])
const('oh_vals', [0.0, 1.0])
const('ch2_idx', [2], 'i')
const('ch4_idx', [4], 'i')
nonbg = nd('Slice', ['input', 'idx_1', 'idx_10', 'idx_1'], [[1, 9, 30, 30]])
n_total = nd('ReduceSum', [nonbg, 'axes123'], [[1, 1, 1, 1]], keepdims=1)
row_plus_col = nd('Add', ['row_grid', 'col_grid'], [[1, 1, 30, 30]])
n_minus_1 = nd('Sub', [n_total, 'c_one'], [[1, 1, 1, 1]])
on_antidiag = nd('Equal', [row_plus_col, n_minus_1], [([1, 1, 30, 30], TensorProto.BOOL)])
col_gt_0 = nd('Greater', ['col_grid', 'c_zero'], [([1, 1, 30, 30], TensorProto.BOOL)])
row_lt_nm1 = nd('Less', ['row_grid', n_minus_1], [([1, 1, 30, 30], TensorProto.BOOL)])
ad_and_col = nd('And', [on_antidiag, col_gt_0], [([1, 1, 30, 30], TensorProto.BOOL)])
antidiag_mask = nd('And', [ad_and_col, row_lt_nm1], [([1, 1, 30, 30], TensorProto.BOOL)])
on_bottom = nd('Equal', ['row_grid', n_minus_1], [([1, 1, 30, 30], TensorProto.BOOL)])
bottom_and_col = nd('And', [on_bottom, col_gt_0], [([1, 1, 30, 30], TensorProto.BOOL)])
col_lt_n_bool = nd('Less', ['col_grid', n_total], [([1, 1, 30, 30], TensorProto.BOOL)])
bottom_mask = nd('And', [bottom_and_col, col_lt_n_bool], [([1, 1, 30, 30], TensorProto.BOOL)])
antidiag_f = nd('Cast', [antidiag_mask], [[1, 1, 30, 30]], to=1)
bottom_f = nd('Cast', [bottom_mask], [[1, 1, 30, 30]], to=1)
col_eq_0 = nd('Equal', ['col_grid', 'c_zero'], [([1, 1, 30, 30], TensorProto.BOOL)])
col_eq_0_f = nd('Cast', [col_eq_0], [[1, 1, 30, 30]], to=1)
row_lt_n = nd('Less', ['row_grid', n_total], [([1, 1, 30, 30], TensorProto.BOOL)])
row_lt_n_f = nd('Cast', [row_lt_n], [[1, 1, 30, 30]], to=1)
col_lt_n = nd('Less', ['col_grid', n_total], [([1, 1, 30, 30], TensorProto.BOOL)])
col_lt_n_f = nd('Cast', [col_lt_n], [[1, 1, 30, 30]], to=1)
inside_grid_f = nd('Mul', [row_lt_n_f, col_lt_n_f], [[1, 1, 30, 30]])
col0_mask = nd('Mul', [col_eq_0_f, inside_grid_f], [[1, 1, 30, 30]])
col0_values = nd('Mul', ['input', col0_mask], [[1, 10, 30, 30]])
ch2_onehot = nd('OneHot', ['ch2_idx', 'depth_10', 'oh_vals'], [[1, 10]], axis=1)
ch2_onehot_4d = nd('Reshape', [ch2_onehot, 'shape_1_10_1_1'], [[1, 10, 1, 1]])
antidiag_ch2 = nd('Mul', [ch2_onehot_4d, antidiag_f], [[1, 10, 30, 30]])
ch4_onehot = nd('OneHot', ['ch4_idx', 'depth_10', 'oh_vals'], [[1, 10]], axis=1)
ch4_onehot_4d = nd('Reshape', [ch4_onehot, 'shape_1_10_1_1'], [[1, 10, 1, 1]])
bottom_ch4 = nd('Mul', [ch4_onehot_4d, bottom_f], [[1, 10, 30, 30]])
occupied_1 = nd('Add', [col0_mask, antidiag_f], [[1, 1, 30, 30]])
occupied = nd('Add', [occupied_1, bottom_f], [[1, 1, 30, 30]])
ch0_mask_raw = nd('Sub', [inside_grid_f, occupied], [[1, 1, 30, 30]])
const('ch0_idx', [0], 'i')
ch0_onehot = nd('OneHot', ['ch0_idx', 'depth_10', 'oh_vals'], [[1, 10]], axis=1)
ch0_onehot_4d = nd('Reshape', [ch0_onehot, 'shape_1_10_1_1'], [[1, 10, 1, 1]])
bg_ch0 = nd('Mul', [ch0_onehot_4d, ch0_mask_raw], [[1, 10, 30, 30]])
out_1 = nd('Add', [col0_values, antidiag_ch2], [[1, 10, 30, 30]])
out_2 = nd('Add', [out_1, bottom_ch4], [[1, 10, 30, 30]])
final = nd('Add', [out_2, bg_ch0], [[1, 10, 30, 30]])
return b.finish('task084', last_tensor=final)
if __name__ == '__main__':
build_and_validate(build_task084, task_num=84)
|