def test_trt_int8_calibration_cache(self):
with tempfile.NamedTemporaryFile() as outpath:
run_polygraphy_run([
ONNX_MODELS["identity"].path, "--trt", "--int8",
"--calibration-cache", outpath.name
])
check_file_non_empty(outpath.name)
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_engine_sanity(self, run_inspect_model):
with tempfile.NamedTemporaryFile() as outpath:
run_polygraphy_run([
ONNX_MODELS["identity"].path, "--trt", "--save-engine",
outpath.name
])
run_inspect_model([outpath.name, "--model-type=engine"])
def test_trt_multiple_profiles(self):
run_polygraphy_run([
ONNX_MODELS["dynamic_identity"].path, "--trt", "--onnxrt",
"--trt-min-shapes", "X,1x2x1x1", "--trt-opt-shapes", "X,1x2x1x1",
"--trt-max-shapes", "X,1x2x1x1", "--trt-min-shapes", "X,1x2x4x4",
"--trt-opt-shapes", "X,1x2x4x4", "--trt-max-shapes", "X,1x2x4x4"
])
def test_custom_per_output_tolerance(self):
run_polygraphy_run([
ONNX_MODELS["identity_identity"].path, "--onnxrt", "--onnxrt",
"--onnx-outputs", "mark", "all", "--atol", "identity_out_0,1.0",
"identity_out_2,3.0", "0.5", "--rtol", "identity_out_0,1.0",
"identity_out_2,3.0", "0.5"
])
def test_onnx_rt_exclude_outputs_with_layerwise(self):
with tempfile.NamedTemporaryFile() as outfile0:
run_polygraphy_run([ONNX_MODELS["identity_identity"].path, "--onnxrt", "--onnx-outputs", "mark", "all", "--onnx-exclude-outputs", "identity_out_2", "--save-outputs", outfile0.name])
results = load_json(outfile0.name)
[result] = list(results.values())[0]
assert len(result) == 1
assert "identity_out_0" in result
def test_int8_calibration_cache(self):
with tempfile.NamedTemporaryFile() as outpath:
cmd = [ONNX_MODELS["identity"].path, "--trt", "--int8", "--calibration-cache", outpath.name]
if mod.version(trt.__version__) >= mod.version("7.0"):
cmd += ["--onnxrt"]
run_polygraphy_run(cmd)
check_file_non_empty(outpath.name)
def identity_engine():
with tempfile.NamedTemporaryFile() as outpath:
run_polygraphy_run([
ONNX_MODELS["identity"].path, "--trt", "--save-engine",
outpath.name
])
yield outpath.name
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_tf_save_pb(self):
with tempfile.NamedTemporaryFile() as outpath:
run_polygraphy_run([
TF_MODELS["identity"].path, "--tf",
"--gpu-memory-fraction=0.5", "--save-pb", outpath.name
])
check_file_non_empty(outpath.name)
def test_save_load_engine(self):
with tempfile.NamedTemporaryFile() as outpath:
run_polygraphy_run([
ONNX_MODELS["identity"].path, "--trt", "--save-engine",
outpath.name
])
assert is_file_non_empty(outpath.name)
run_polygraphy_run(["--trt", outpath.name, "--model-type=engine"])
def test_tf_save_tensorboard(self):
with tempfile.TemporaryDirectory() as outdir:
run_polygraphy_run([
TF_MODELS["identity"].path, "--tf",
"--gpu-memory-fraction=0.5", "--save-tensorboard", outdir
])
files = glob.glob("{:}{:}*".format(outdir, os.path.sep))
assert len(files) == 1
def test_tf2onnx_save_onnx(self):
with util.NamedTemporaryFile() as outpath:
run_polygraphy_run([
TF_MODELS["identity"].path, "--onnxrt", "--model-type=frozen",
"--save-onnx", outpath.name
])
assert is_file_non_empty(outpath.name)
assert onnx.load(outpath.name)
def test_int8_calibration_base_class(self, base_class):
cmd = [
ONNX_MODELS["identity"].path, "--trt", "--int8",
"--calibration-base-class", base_class
]
if mod.version(trt.__version__) >= mod.version("7.0"):
cmd += ["--onnxrt"]
run_polygraphy_run()
def test_polygraphy_run_gen_script(self):
with tempfile.NamedTemporaryFile(mode="w") as f:
run_polygraphy_run(["--gen-script={:}".format(f.name), ONNX_MODELS["identity"].path])
with open(f.name, "r") as script:
print(script.read())
env = copy.deepcopy(os.environ)
env.update({"PYTHONPATH": ROOT_DIR})
check_subprocess(sp.run([sys.executable, f.name], env=env))
def test_plugins(self):
run_polygraphy_run([
ONNX_MODELS["identity"].path,
"--trt",
"--plugins",
"nvinfer_plugin.dll"
if sys.platform.startswith("win") else "libnvinfer_plugin.so",
])
def test_onnx_rt_save_onnx(self):
with tempfile.NamedTemporaryFile() as outpath:
run_polygraphy_run([
ONNX_MODELS["identity"].path, "--onnxrt", "--save-onnx",
outpath.name
])
assert is_file_non_empty(outpath.name)
assert onnx.load(outpath.name)
def test_tf_save_timeline(self):
with tempfile.NamedTemporaryFile() as outpath:
run_polygraphy_run([
TF_MODELS["identity"].path, "--tf",
"--gpu-memory-fraction=0.5", "--save-timeline", outpath.name
])
timelines = glob.glob(os.path.join(outpath.name, "*"))
for timeline in timelines:
assert is_file_non_empty(timeline)
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_precision_worst_first_sanity():
with tempfile.NamedTemporaryFile() as outpath:
run_polygraphy_run([
ONNX_MODELS["identity"].path, "--onnxrt", "--save-results",
outpath.name, "--onnx-outputs", "mark", "all"
])
run_polygraphy_precision([
"worst-first", ONNX_MODELS["identity"].path, "--golden",
outpath.name, "--int8", "--trt-outputs", "mark", "all"
])
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_polygraphy_precision_linear_sanity():
with tempfile.NamedTemporaryFile() as outpath:
run_polygraphy_run([
ONNX_MODELS["identity"].path, "--onnxrt", "--save-results",
outpath.name
])
run_polygraphy_precision([
"linear", ONNX_MODELS["identity"].path, "--golden", outpath.name,
"--int8"
])
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_tactic_replay(self):
with tempfile.NamedTemporaryFile() as tactic_replay:
run_polygraphy_run([
ONNX_MODELS["identity"].path, "--trt", "--save-tactics",
tactic_replay.name
])
assert is_file_non_empty(tactic_replay.name)
run_polygraphy_run([
ONNX_MODELS["identity"].path, "--trt", "--load-tactics",
tactic_replay.name
])
def test_trt_layerwise_outputs(self):
with tempfile.NamedTemporaryFile() as outfile0:
run_polygraphy_run([
ONNX_MODELS["identity_identity"].path, "--trt",
"--trt-outputs", "mark", "all", "--save-results", outfile0.name
])
results = misc.pickle_load(outfile0.name)
[result] = list(results.values())[0]
assert len(result) == 2
assert "identity_out_0" in result
assert "identity_out_2" in result
def test_explicit_profile_implicit_runtime_shape(self):
run_polygraphy_run([
ONNX_MODELS["dynamic_identity"].path,
"--trt",
"--onnxrt",
"--trt-min-shapes",
"X:[1,2,1,1]",
"--trt-opt-shapes",
"X:[1,2,1,1]",
"--trt-max-shapes",
"X:[1,2,1,1]",
])
def test_save_load_outputs(self, tmp_path):
OUTFILE0 = os.path.join(tmp_path, "outputs0.pkl")
OUTFILE1 = os.path.join(tmp_path, "outputs1.pkl")
run_polygraphy_run([
ONNX_MODELS["identity"].path, "--onnxrt", "--save-results",
OUTFILE0
])
run_polygraphy_run([
ONNX_MODELS["identity"].path, "--onnxrt", "--save-results",
OUTFILE1
])
status = run_polygraphy_run([
ONNX_MODELS["identity"].path, "--onnxrt", "--load-results",
OUTFILE0, OUTFILE1
])
assert "Difference is within tolerance" in status.stdout + status.stderr # Make sure it actually compared stuff.
# Should work with only one file
status = run_polygraphy_run(
[ONNX_MODELS["identity"].path, "--load-results", OUTFILE0])
assert "Difference is within tolerance" not in status.stdout + status.stderr # Make sure it DIDN'T compare stuff.
# Should work even with no runners specified
status = run_polygraphy_run([
ONNX_MODELS["identity"].path, "--load-results", OUTFILE0, OUTFILE1
])
assert "Difference is within tolerance" in status.stdout + status.stderr # Make sure it actually compared stuff.
# Should work even when comparing a single runner to itself.
status = run_polygraphy_run([
ONNX_MODELS["identity"].path, "--load-results", OUTFILE0, OUTFILE0
])
assert "Difference is within tolerance" in status.stdout + status.stderr # Make sure it actually compared stuff.
def test_explicit_profile_opt_runtime_shapes_differ(self):
run_polygraphy_run([
ONNX_MODELS["dynamic_identity"].path,
"--trt",
"--onnxrt",
"--input-shapes",
"X:[1,2,2,2]",
"--trt-min-shapes",
"X:[1,2,1,1]",
"--trt-opt-shapes",
"X:[1,2,3,3]",
"--trt-max-shapes",
"X:[1,2,4,4]",
])
def test_data_loader_script_calibration(self):
with tempfile.NamedTemporaryFile("w+", suffix=".py") as f:
f.write(
dedent("""
import numpy as np
def load_data():
for _ in range(5):
yield {"x": np.ones((1, 1, 2, 2), dtype=np.float32) * 6.4341}
"""))
f.flush()
run_polygraphy_run([
ONNX_MODELS["identity"].path, "--trt", "--int8",
"--data-loader-script", f.name
])
def test_multiple_profiles(self):
run_polygraphy_run([
ONNX_MODELS["dynamic_identity"].path,
"--trt",
"--onnxrt",
"--trt-min-shapes",
"X:[1,2,1,1]",
"--trt-opt-shapes",
"X:[1,2,1,1]",
"--trt-max-shapes",
"X:[1,2,1,1]",
"--trt-min-shapes",
"X:[1,2,4,4]",
"--trt-opt-shapes",
"X:[1,2,4,4]",
"--trt-max-shapes",
"X:[1,2,4,4]",
])