def test_external_data(self):
with util.NamedTemporaryFile() as path, util.NamedTemporaryFile(
) as data:
model = OnnxFromPath(ONNX_MODELS["const_foldable"].path)
loader = SaveOnnx(model,
path.name,
external_data_path=data.name,
size_threshold=0)
loader()
assert is_file_non_empty(path.name)
assert is_file_non_empty(data.name)
def test_external_data(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=0"]
)
arg_group.save_onnx(model)
assert is_file_non_empty(path.name)
assert is_file_non_empty(data.name)
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
])
assert is_file_non_empty(outpath.name)
def test_timing_cache(self):
with tempfile.TemporaryDirectory() as dir:
# Test with files that haven't already been created instead of using NamedTemporaryFile().
total_cache = os.path.join(dir, "total.cache")
identity_cache = os.path.join(dir, "identity.cache")
run_polygraphy_run([
ONNX_MODELS["const_foldable"].path, "--trt", "--timing-cache",
total_cache
])
assert is_file_non_empty(total_cache)
const_foldable_cache_size = get_file_size(total_cache)
run_polygraphy_run([
ONNX_MODELS["identity"].path, "--trt", "--timing-cache",
identity_cache
])
identity_cache_size = get_file_size(identity_cache)
run_polygraphy_run([
ONNX_MODELS["identity"].path, "--trt", "--timing-cache",
total_cache
])
total_cache_size = get_file_size(total_cache)
# The total cache should be larger than either of the individual caches.
assert total_cache_size > const_foldable_cache_size and total_cache_size > identity_cache_size
# The total cache should also be smaller than or equal to the sum of the individual caches since
# header information should not be duplicated.
assert total_cache_size <= (const_foldable_cache_size +
identity_cache_size)
def test_save_onnx(self):
with tempfile.TemporaryDirectory() as outdir:
outpath = os.path.join(outdir, "test", "nested")
loader = SaveOnnx(OnnxFromPath(ONNX_MODELS["identity"].path),
path=outpath)
loader()
assert is_file_non_empty(outpath)
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_save_timeline(self):
model = TF_MODELS["identity"]
with util.NamedTemporaryFile() as outpath:
with TfRunner(SessionFromGraph(model.loader),
allow_growth=True,
save_timeline=outpath.name) as runner:
model.check_runner(runner)
assert is_file_non_empty(outpath.name)
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_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_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)
assert is_file_non_empty(outpath.name)
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_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_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_pb(self):
with util.NamedTemporaryFile() as outpath:
tf_loader = SaveGraph(GraphFromFrozen(TF_MODELS["identity"].path),
path=outpath.name)
tf_loader()
assert is_file_non_empty(outpath.name)
def test_save_engine(self, identity_network):
with tempfile.NamedTemporaryFile() as outpath:
engine_loader = SaveEngine(EngineFromNetwork(identity_network),
path=outpath.name)
with engine_loader():
assert is_file_non_empty(outpath.name)
def test_save_onnx(self):
with tempfile.NamedTemporaryFile() as outpath:
loader = SaveOnnx(OnnxFromPath(ONNX_MODELS["identity"].path),
path=outpath.name)
loader()
assert is_file_non_empty(outpath.name)
