File size: 2,038 Bytes
24e7fa4 | 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 | """plan() on tiny synthetic graphs (no real weights)."""
from __future__ import annotations
import re
import unittest
import numpy as np
import onnx
from onnx import TensorProto, helper, numpy_helper
from tools.maya_export.quant_depth_frame import plan
def _graph(consumers):
"""One fp16 (4, 3) weight 'w' with one Cast per entry of consumers ('mm' = Transpose->MatMul, 'bad' = Add)."""
w = numpy_helper.from_array(np.ones((4, 3), np.float16), "w")
nodes, outs = [], []
for i, kind in enumerate(consumers):
nodes.append(helper.make_node("Cast", ["w"], [f"wf{i}"], name=f"cast{i}", to=TensorProto.FLOAT))
if kind == "mm":
nodes.append(helper.make_node("Transpose", [f"wf{i}"], [f"wt{i}"], name=f"tr{i}", perm=[1, 0]))
nodes.append(helper.make_node("MatMul", ["x", f"wt{i}"], [f"y{i}"], name=f"mm{i}"))
else:
nodes.append(helper.make_node("Add", [f"wf{i}", f"wf{i}"], [f"y{i}"], name=f"add{i}"))
outs.append(helper.make_tensor_value_info(f"y{i}", TensorProto.FLOAT, None))
x = helper.make_tensor_value_info("x", TensorProto.FLOAT, [2, 3])
return helper.make_model(helper.make_graph(nodes, "g", [x], outs, initializer=[w]))
class TestPlan(unittest.TestCase):
def test_single_cast(self):
targets, drop = plan(_graph(["mm"]), None)
self.assertEqual([k for k, _ in targets["w"]], ["matmul"])
self.assertEqual(drop, {"cast0", "tr0"})
def test_shared_weight_accumulates_all_uses(self):
targets, drop = plan(_graph(["mm", "mm"]), None)
self.assertEqual([n.name for _, n in targets["w"]], ["mm0", "mm1"])
self.assertEqual(drop, {"cast0", "tr0", "cast1", "tr1"})
def test_mixed_convertible_and_not_raises(self):
with self.assertRaises(ValueError):
plan(_graph(["mm", "bad"]), None)
def test_keep_regex_skips(self):
targets, _ = plan(_graph(["mm"]), re.compile("^w$"))
self.assertEqual(targets, {})
if __name__ == "__main__":
unittest.main()
|