def test_inputs(self, opts):
with tempfile.NamedTemporaryFile() as outpath:
run_polygraphy_run([
ONNX_MODELS["identity"].path, "--onnxrt", "--save-inputs",
outpath.name
])
run_polygraphy_inspect(["data", outpath.name] + opts)
def test_model_trt_network_script(self):
script = dedent("""
from polygraphy.backend.trt import CreateNetwork
from polygraphy import func
import tensorrt as trt
@func.extend(CreateNetwork())
def load_network(builder, network):
inp = network.add_input("input", dtype=trt.float32, shape=(1, 1))
out = network.add_identity(inp).get_output(0)
network.mark_output(out)
""")
with tempfile.NamedTemporaryFile("w+", suffix=".py") as f:
f.write(script)
f.flush()
run_polygraphy_inspect(["model", f.name])
def test_model_onnx(self, case):
model, mode, expected = case
status = run_polygraphy_inspect(
["model", ONNX_MODELS[model].path, "--mode={:}".format(mode)], disable_verbose=True
)
expected = dedent(expected).strip()
actual = "\n".join(status.stdout.splitlines()[1:]) # Ignore loading message
check_lines_match(actual, expected)
def test_capability(self, case):
model, expected_files, expected_summary = case
with tempfile.TemporaryDirectory() as outdir:
status = run_polygraphy_inspect(
["capability", ONNX_MODELS[model].path, "-o", os.path.join(outdir, "subgraphs")],
)
assert sorted(map(os.path.basename, glob.glob(os.path.join(outdir, "subgraphs", "**")))) == sorted(
expected_files
)
assert dedent(expected_summary).strip() in status.stdout
def test_show_tactics(self, case):
with util.NamedTemporaryFile() as replay:
model_name, expected = case
run_polygraphy_run([ONNX_MODELS[model_name].path, "--trt", "--save-tactics", replay.name])
status = run_polygraphy_inspect(["tactics", replay.name], disable_verbose=True)
expected = dedent(expected).strip()
actual = status.stdout
check_lines_match(actual, expected, should_check_line=lambda line: "Algorithm: " not in line)
def test_polygraphy_inspect_model_onnx(run_inspect_model, case):
model, mode, expected = case
status = run_polygraphy_inspect(["model", ONNX_MODELS[model].path, "--mode={:}".format(mode)], disable_verbose=True)
expected = dedent(expected).strip()
actual = status.stdout.decode()
print("Actual output:\n{:}".format(actual))
for acline, exline in zip(actual.splitlines(), expected.splitlines()):
acline = acline.rstrip()
exline = exline.rstrip()
print("Checking line : {:}".format(acline))
print("Expecting line: {:}".format(exline))
assert acline == exline
def run_inspect_model(request):
yield lambda additional_opts: run_polygraphy_inspect(
["model"] + ["--mode={:}".format(request.param)] + additional_opts)
