def test_can_isolate_node(self, fail_node, mode):
with tempfile.TemporaryDirectory() as outdir:
run_polygraphy_debug(["reduce", ONNX_MODELS["reducable"].path, "--output=reduced.onnx", "--mode", mode, "--show-output",
"--min-good=good_reduced.onnx",
"--check", TestReduce.FAKE_REDUCE_CHECKER, "polygraphy_debug.onnx", "--fail-node", fail_node],
disable_verbose=True, cwd=outdir)
model = onnx_from_path(os.path.join(outdir, "reduced.onnx"))
min_good_path = os.path.join(outdir, "good_reduced.onnx")
good_model = None
if os.path.exists(min_good_path):
good_model = onnx_from_path(min_good_path)
# The model should only contain one node - the failing one.
# One exception - since bisect depends on node ordering, it sometimes doesn't
# reduce branches to the maximum possible extent.
if mode == "bisect" and fail_node == "onnx_graphsurgeon_node_1":
assert len(model.graph.node) == 3
elif mode == "bisect" and fail_node == "onnx_graphsurgeon_node_7":
assert len(model.graph.node) == 2
else:
assert len(model.graph.node) == 1
node_names = [node.name for node in model.graph.node]
assert fail_node in node_names
# For now we're just doing a very basic sanity check for --min-good
if good_model:
assert model != good_model
def test_no_reduce_outputs(self):
with tempfile.TemporaryDirectory() as outdir:
run_polygraphy_debug(
[
"reduce",
ONNX_MODELS["reducable"].path,
"--output=reduced.onnx",
"--show-output",
"--no-reduce-outputs",
"--mode=linear",
"--check",
TestReduce.FAKE_REDUCE_CHECKER,
"polygraphy_debug.onnx",
"--fail-node",
"onnx_graphsurgeon_node_3",
],
disable_verbose=True,
cwd=outdir,
)
model = onnx_from_path(os.path.join(outdir, "reduced.onnx"))
assert len(model.graph.node) == 4
assert len(model.graph.input) == 1
assert len(model.graph.output) == 2
assert model.graph.input[0].name == "Y0"
node_names = [node.name for node in model.graph.node]
assert "onnx_graphsurgeon_node_7" in node_names
def test_reduce_custom_fail_message(self, fail_code_arg):
with tempfile.TemporaryDirectory() as outdir:
# fake_reduce_checker will alternate error messages based on whether an arbitrary node is present in the model.
run_polygraphy_debug(
[
"reduce",
ONNX_MODELS["reducable"].path,
"--output=reduced.onnx",
"--show-output",
"--fail-regex",
"REALLY BAD",
"BAD NODE",
] + fail_code_arg + [
"--check",
TestReduce.FAKE_REDUCE_CHECKER,
"polygraphy_debug.onnx",
"--fail-node",
"onnx_graphsurgeon_node_5",
"--fail-return-code=0",
],
disable_verbose=True,
cwd=outdir,
)
model = onnx_from_path(os.path.join(outdir, "reduced.onnx"))
assert len(model.graph.node) == 1
assert model.graph.node[0].name == "onnx_graphsurgeon_node_5"
def test_reduce_with_constant(self):
# Should be no failure when models including Constant nodes use fallback
# shape inference; Constant nodes will be lowered to constant tensors.
with tempfile.TemporaryDirectory() as outdir:
run_polygraphy_debug(
[
"reduce",
ONNX_MODELS["reducable_with_const"].path,
"--no-shape-inference",
"--mode=linear",
"--output=reduced.onnx",
] + [
"--check",
TestReduce.FAKE_REDUCE_CHECKER,
"polygraphy_debug.onnx",
"--fail-node",
"onnx_graphsurgeon_node_3",
],
disable_verbose=True,
cwd=outdir,
)
model = onnx_from_path(os.path.join(outdir, "reduced.onnx"))
graph = gs.import_onnx(model)
assert len(graph.nodes) == 1
assert graph.nodes[0].name == "onnx_graphsurgeon_node_3"
# Outputs of Constant nodes should not become Variables; thus the model should have no inputs.
assert not graph.inputs
def test_layerwise(self, copy):
original_model = onnx_from_path(ONNX_MODELS["identity_identity"].path)
loader = ModifyOutputs(original_model,
outputs=constants.MARK_ALL,
copy=copy)
model = loader()
assert len(original_model.graph.output) == 1 or not copy
assert len(model.graph.output) == 2
def test_basic(self, copy):
original_model = onnx_from_path(ONNX_MODELS["identity_identity"].path)
loader = ConvertToFp16(original_model, copy=copy)
model = loader()
assert original_model.graph.input[
0].type.tensor_type.elem_type == 1 or not copy
assert model.graph.value_info[0].type.tensor_type.elem_type == 10
def extract_model():
input_metadata = TensorMetadata().add("X",
dtype=np.float32,
shape=(64, 64))
output_metadata = TensorMetadata().add("identity_out_0",
dtype=np.float32,
shape=None)
return onnx_from_path(
ONNX_MODELS["identity_identity"].path), input_metadata, output_metadata
def test_basic(self, partitioning, fold_shapes, copy):
original_model = onnx_from_path(ONNX_MODELS["const_foldable"].path)
loader = FoldConstants(original_model,
partitioning=partitioning,
fold_shapes=fold_shapes,
copy=copy)
model = loader()
assert len(original_model.graph.node) != 1 or not copy
assert len(model.graph.node) == 1
def test_external_data(self):
model = onnx_from_path(ONNX_MODELS["const_foldable"].path)
arg_group = ArgGroupTestHelper(OnnxSaveArgs(), deps=[ModelArgs(), OnnxLoaderArgs()])
with tempfile.NamedTemporaryFile() as path, tempfile.NamedTemporaryFile() as data:
arg_group.parse_args(["-o", path.name, "--save-external-data", data.name])
arg_group.save_onnx(model)
check_file_non_empty(path.name)
check_file_non_empty(data.name)
def test_reduce_custom_return_code(self):
with tempfile.TemporaryDirectory() as outdir:
run_polygraphy_debug(["reduce", ONNX_MODELS["reducable"].path, "--output=reduced.onnx", "--show-output", "--fail-code=1", # Only 1s are real failures.
"--check", TestReduce.FAKE_REDUCE_CHECKER, "polygraphy_debug.onnx", "--fail-node", "onnx_graphsurgeon_node_5", "--default-return-code=2"],
disable_verbose=True, cwd=outdir)
model = onnx_from_path(os.path.join(outdir, "reduced.onnx"))
assert len(model.graph.node) == 1
assert model.graph.node[0].name == "onnx_graphsurgeon_node_5"
def test_reduce_shape_inference(self, opts):
with tempfile.TemporaryDirectory() as outdir:
status = run_polygraphy_debug(["reduce", ONNX_MODELS["dynamic_identity"].path, "--output=reduced.onnx",
"--show-output", "--model-input-shapes=X:[1,2,5,5]"] + opts
+ ["--check", "false"],
disable_verbose=True, cwd=outdir)
model = onnx_from_path(os.path.join(outdir, "reduced.onnx"))
graph = gs.import_onnx(model)
assert tuple(graph.inputs[0].shape) == (1, 2, 5, 5)
assert tuple(graph.outputs[0].shape) == (1, 2, 5, 5)
def test_size_threshold(self):
model = onnx_from_path(ONNX_MODELS["const_foldable"].path)
arg_group = ArgGroupTestHelper(OnnxSaveArgs(), deps=[ModelArgs(), OnnxLoaderArgs()])
with util.NamedTemporaryFile() as path, util.NamedTemporaryFile() as data:
arg_group.parse_args(
["-o", path.name, "--save-external-data", data.name, "--external-data-size-threshold=1024"]
)
arg_group.save_onnx(model)
assert is_file_non_empty(path.name)
assert is_file_empty(data.name)
def test_no_reduce_required_branches(self, fail_nodes):
with tempfile.TemporaryDirectory() as outdir:
run_polygraphy_debug(["reduce", ONNX_MODELS["reducable"].path, "--output=reduced.onnx", "--show-output",
"--check", TestReduce.FAKE_REDUCE_CHECKER, "polygraphy_debug.onnx",
"--fail-node"] + fail_nodes,
disable_verbose=True, cwd=outdir)
model = onnx_from_path(os.path.join(outdir, "reduced.onnx"))
node_names = [node.name for node in model.graph.node]
assert all(fail_node in node_names for fail_node in fail_nodes)
assert len(model.graph.node) <= 3 # The branch on the opposite side of the model should be removed.
def test_no_all_tensors_to_one_file(self):
model = onnx_from_path(ONNX_MODELS["const_foldable"].path)
arg_group = ArgGroupTestHelper(OnnxSaveArgs(), deps=[ModelArgs(), OnnxLoaderArgs()])
with tempfile.TemporaryDirectory() as outdir:
path = os.path.join(outdir, "model.onnx")
arg_group.parse_args(
[
"-o",
path,
"--save-external-data",
"--external-data-size-threshold=0",
"--no-save-all-tensors-to-one-file",
]
)
arg_group.save_onnx(model)
assert is_file_non_empty(path)
outfiles = glob.glob(os.path.join(outdir, "*"))
assert len(outfiles) == 4
def test_save_to_disk_on_size_threshold(self):
model = onnx_from_path(ONNX_MODELS["const_foldable"].path)
model = infer_shapes(model, save_to_disk_threshold_bytes=0)
self.check_model(model)
def test_get_num_nodes():
model = onnx_from_path(ONNX_MODELS["scan"].path)
assert onnx_util.get_num_nodes(model) == 3 # Should count subgraph nodes.
def test_model(self):
original_model = onnx_from_path(ONNX_MODELS["identity_identity"].path)
model = infer_shapes(original_model)
self.check_model(model)
