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)