def test_polygraphy_surgeon_extract_sanity():
with tempfile.NamedTemporaryFile() as modelpath:
run_polygraphy_surgeon([
"extract", ONNX_MODELS["identity_identity"].path, "-o",
modelpath.name, "--inputs", "identity_out_0,auto,auto"
])
run_polygraphy_run([modelpath.name, "--model-type=onnx", "--onnxrt"])
def test_sanity_dim_param(self):
with tempfile.NamedTemporaryFile() as modelpath:
run_polygraphy_surgeon([
"extract", ONNX_MODELS["dim_param"].path, "-o", modelpath.name
])
run_polygraphy_run(
[modelpath.name, "--model-type=onnx", "--onnxrt"])
def test_with_attributes(self):
with tempfile.NamedTemporaryFile() as outmodel:
# str_attr='0' should be interpreted as a string, not an int
# float_attr=0.0 should be interpreted as a float, not an int
# int_attr=0 should be interpreted as an int
run_polygraphy_surgeon(["insert", ONNX_MODELS["identity_identity"].path, "-o", outmodel.name, "--inputs=X",
"--outputs=X", "--op=FakeOp",
"--attrs",
"str_attr='0'", "int_attr=0", "float_attr=0.0", "other_str_attr=name",
"str_list_attr=['0','1']", "int_list_attr=[1,2,3]", "float_list_attr=[0.0,-1.0,-2.0]"])
model = self.check_insert_model(outmodel.name, ["FakeOp", "Identity", "Identity"], ["X"], ["identity_out_2"])
node = model.graph.node[0]
attrs = node.attribute
assert attrs[0].name == "str_attr"
assert attrs[0].s == b"0"
assert attrs[1].name == "int_attr"
assert attrs[1].i == 0
assert attrs[2].name == "float_attr"
assert attrs[2].f == 0.0
assert attrs[3].name == "other_str_attr"
assert attrs[3].s == b"name"
assert attrs[4].name == "str_list_attr"
assert attrs[4].strings == [b"0", b"1"]
assert attrs[5].name == "int_list_attr"
assert attrs[5].ints == [1, 2, 3]
assert attrs[6].name == "float_list_attr"
assert attrs[6].floats == [0.0, -1.0, -2.0]
def test_cleanup(self):
with tempfile.NamedTemporaryFile(suffix=".onnx") as outmodel:
run_polygraphy_surgeon(["sanitize", ONNX_MODELS["identity_identity"].path, "-o", outmodel.name, "--outputs", "identity_out_0", "--cleanup"])
model = onnx.load(outmodel.name)
assert len(model.graph.node) == 1
assert model.graph.output[0].name == "identity_out_0"
def test_multi_use_input(self):
with tempfile.NamedTemporaryFile() as outmodel:
run_polygraphy_surgeon(["insert", ONNX_MODELS["reducable"].path, "-o", outmodel.name, "--inputs=add_out_4",
"--outputs=identity_out_8", "--op=FakeOp"])
model = self.check_insert_model(outmodel.name, ["Identity", "Identity", "Add", "FakeOp", "Identity"], ["X0", "Y0"], ["identity_out_6", "identity_out_8"])
other_branch_node = model.graph.node[-1]
assert other_branch_node.name == "onnx_graphsurgeon_node_7"
assert other_branch_node.input == ["add_out_4"]
def test_modify_onnx_outputs(self):
with tempfile.NamedTemporaryFile(suffix=".onnx") as outmodel:
run_polygraphy_surgeon(
["sanitize", ONNX_MODELS["identity_identity"].path, "-o", outmodel.name, "--outputs", "mark", "all"]
)
model = onnx.load(outmodel.name)
assert len(model.graph.output) == 2
def test_at_graph_input(self):
with tempfile.NamedTemporaryFile() as modelpath:
run_polygraphy_surgeon([
"insert", ONNX_MODELS["identity_identity"].path, "-o",
modelpath.name, "--inputs=X", "--outputs=X", "--op=FakeOp"
])
self.check_insert_model(modelpath.name,
["FakeOp", "Identity", "Identity"], ["X"],
["identity_out_2"])
def test_override_shapes_no_clear_const_tensors_meta(self):
with util.NamedTemporaryFile() as outmodel:
run_polygraphy_surgeon([
"sanitize",
ONNX_MODELS["const_foldable"].path,
"-o",
outmodel.name,
"--override-input-shapes=input:[1,3]",
])
def test_sanity(self):
with tempfile.NamedTemporaryFile() as modelpath:
run_polygraphy_surgeon([
"sanitize", ONNX_MODELS["const_foldable"].path, "-o",
modelpath.name, "--fold-constants"
])
model = onnx.load(modelpath.name)
assert len(model.graph.node) == 1
def test_polygraphy_surgeon_extract_fallback_shape_inference():
with tempfile.NamedTemporaryFile() as modelpath:
# Force fallback shape inference by disabling ONNX shape inference
run_polygraphy_surgeon([
"extract", ONNX_MODELS["identity_identity"].path, "-o",
modelpath.name, "--inputs", "identity_out_0,auto,auto",
"--outputs", "identity_out_2,auto", "--no-shape-inference"
])
run_polygraphy_run([modelpath.name, "--model-type=onnx", "--onnxrt"])
def test_polygraphy_surgeon_sanity():
with tempfile.NamedTemporaryFile(
) as configpath, tempfile.NamedTemporaryFile() as modelpath:
run_polygraphy_surgeon(
["prepare", ONNX_MODELS["identity"].path, "-o", configpath.name])
run_polygraphy_surgeon([
"operate", ONNX_MODELS["identity"].path, "-c", configpath.name,
"-o", modelpath.name
])
run_polygraphy_run([modelpath.name, "--model-type=onnx", "--onnxrt"])
def test_force_fallback_shape_inference_will_override_model_shapes(self):
with tempfile.NamedTemporaryFile() as outmodel:
run_polygraphy_surgeon(["extract", ONNX_MODELS["dynamic_identity"].path, "-o", outmodel.name, "--outputs", "Y:auto", "--force-fallback-shape-inference"])
onnx_model_sanity_check(outmodel.name)
graph = gs.import_onnx(onnx.load(outmodel.name))
# Inputs should become fixed since fallback shape inference is being forced.
for tensor in graph.tensors().values():
assert tensor.shape is not None
assert tuple(graph.inputs[0].shape) == (1, 2, 1, 1)
assert tuple(graph.outputs[0].shape) == (1, 2, 1, 1)
def test_graph_output(self):
# FakeOp output tensor should be marked as a graph output since identity_out_2 was.
with tempfile.NamedTemporaryFile() as modelpath:
run_polygraphy_surgeon([
"insert", ONNX_MODELS["identity_identity"].path, "-o",
modelpath.name, "--inputs=identity_out_2",
"--outputs=identity_out_2", "--op=FakeOp"
])
self.check_insert_model(
modelpath.name, ["Identity", "Identity", "FakeOp"], ["X"],
["identity_out_2_polygraphy_surgeon_insert_output"])
def test_fold_constants(self, partitioning, fold_shapes):
with tempfile.NamedTemporaryFile() as outmodel:
cmd = ["sanitize", ONNX_MODELS["const_foldable"].path, "-o", outmodel.name, "--fold-constants"]
if fold_shapes:
cmd += [fold_shapes]
if partitioning:
cmd += ["--partitioning", partitioning]
run_polygraphy_surgeon(cmd)
onnx_model_sanity_check(outmodel.name)
model = onnx.load(outmodel.name)
assert len(model.graph.node) == 1
def test_override_shapes_partial_inputs(self):
with util.NamedTemporaryFile() as outmodel:
run_polygraphy_surgeon([
"sanitize",
ONNX_MODELS["dynamic_identity"].path,
"-o",
outmodel.name,
"--override-input-shapes=Y:[1,2,3,4]",
])
model = onnx.load(outmodel.name)
assert model.graph.input[0].type.tensor_type.shape.dim[
2].dim_param == "height"
assert model.graph.input[0].type.tensor_type.shape.dim[
3].dim_param == "width"
def test_override_shapes(self, new_dim):
with tempfile.NamedTemporaryFile() as outmodel:
cmd = ["sanitize", ONNX_MODELS["dynamic_identity"].path, "-o", outmodel.name, "--override-input-shapes=X:[1,2,{new_dim},{new_dim}]".format(new_dim=new_dim)]
run_polygraphy_surgeon(cmd)
onnx_model_sanity_check(outmodel.name)
model = onnx.load(outmodel.name)
shape = []
for dim in model.graph.input[0].type.tensor_type.shape.dim:
assert isinstance(dim.dim_value, int) and dim.dim_value >= 0
shape.append(dim.dim_value)
assert shape == [1, 2, new_dim, new_dim]
def test_override_shapes_no_reorder(self):
with util.NamedTemporaryFile() as outmodel:
run_polygraphy_surgeon([
"sanitize",
ONNX_MODELS["reducable"].path,
"-o",
outmodel.name,
"--override-input-shapes",
"Y0:[5]",
"X0:[5]",
])
model = onnx.load(outmodel.name)
assert model.graph.input[0].name == "X0"
assert model.graph.input[1].name == "Y0"
def test_insert_at_tensor(self):
# Insert a new node in between existing nodes without replacing any existing nodes.
with util.NamedTemporaryFile() as outmodel:
run_polygraphy_surgeon([
"insert",
ONNX_MODELS["identity_identity"].path,
"-o",
outmodel.name,
"--inputs=identity_out_0",
"--outputs=identity_out_0",
"--op=FakeOp",
])
self.check_insert_model(outmodel.name,
["Identity", "FakeOp", "Identity"], ["X"],
["identity_out_2"])
def test_graph_output(self):
# FakeOp output tensor should be marked as a graph output. Name should be preserved - identity_out_2
with util.NamedTemporaryFile() as outmodel:
run_polygraphy_surgeon([
"insert",
ONNX_MODELS["identity_identity"].path,
"-o",
outmodel.name,
"--inputs=identity_out_2",
"--outputs=identity_out_2",
"--op=FakeOp",
])
self.check_insert_model(outmodel.name,
["Identity", "Identity", "FakeOp"], ["X"],
["identity_out_2"])
def test_no_shape_inference_if_has_metadata(self):
with tempfile.NamedTemporaryFile() as outmodel:
status = run_polygraphy_surgeon(
["extract", ONNX_MODELS["identity_identity"].path, "-o", outmodel.name, "--inputs", "X:auto:auto"]
)
onnx_model_sanity_check(outmodel.name)
assert not was_shape_inference_run(status)
def test_fold_constants_single_pass(self):
with tempfile.NamedTemporaryFile() as outmodel:
status = run_polygraphy_surgeon(["sanitize", ONNX_MODELS["const_foldable"].path, "-o", outmodel.name, "--fold-constants", "--num-passes=1"])
assert "Pass 1" in status.stdout
assert "Pass 2" not in status.stdout
onnx_model_sanity_check(outmodel.name)
model = onnx.load(outmodel.name)
assert len(model.graph.node) == 1
def test_fallback_shape_inference_no_onnx_shape_inference(self):
with util.NamedTemporaryFile() as outmodel:
status = run_polygraphy_surgeon([
"extract",
ONNX_MODELS["identity_identity"].path,
"-o",
outmodel.name,
"--inputs",
"identity_out_0:auto:auto",
"--outputs",
"identity_out_2:auto",
"--force-fallback-shape-inference",
])
onnx_model_sanity_check(outmodel.name)
assert not was_shape_inference_run(status)
def test_external_data(self):
with tempfile.TemporaryDirectory() as outdir:
model = ONNX_MODELS["ext_weights"]
outmodel = os.path.join(outdir, "out_model.onnx")
outdata = "ext_weights.data"
assert run_polygraphy_surgeon([
"sanitize",
model.path,
"--external-data-dir",
model.ext_data,
"--fold-constants",
"-o",
outmodel,
"--save-external-data",
outdata,
"--external-data-size-threshold=0",
"-vvvvv",
])
assert is_file_non_empty(outmodel)
assert is_file_non_empty(os.path.join(outdir, outdata))
assert run_polygraphy_run(
[outmodel, "--onnxrt", "--external-data-dir", outdir])
def test_sanity_dim_param(self):
with util.NamedTemporaryFile() as outmodel:
run_polygraphy_surgeon([
"extract", ONNX_MODELS["dim_param"].path, "-o", outmodel.name
])
onnx_model_sanity_check(outmodel.name)
def test_external_data(self):
with tempfile.NamedTemporaryFile(suffix=".onnx") as outmodel, tempfile.NamedTemporaryFile() as data:
model = ONNX_MODELS["ext_weights"]
assert run_polygraphy_surgeon(["sanitize", model.path, "-o", outmodel.name, "--load-external-data", model.ext_data, "--save-external-data", data.name, "-vvvvv"])
check_file_non_empty(outmodel.name)
check_file_non_empty(data.name)
