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 test_save_timeline(self):
model = TF_MODELS["identity"]
with tempfile.NamedTemporaryFile() as outpath:
with TfRunner(SessionFromGraph(model.loader),
allow_growth=True,
save_timeline=outpath.name) as runner:
model.check_runner(runner)
check_file_non_empty(outpath.name)
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_calibrator_with_path_name_cache(self, identity_builder_network):
builder, network = identity_builder_network
data = [{"x": np.ones((1, 1, 2, 2), dtype=np.float32)}]
with tempfile.NamedTemporaryFile() as cache:
calibrator = Calibrator(data, cache=cache.name)
create_config = CreateConfig(int8=True, calibrator=calibrator)
with engine_from_network((builder, network), create_config):
check_file_non_empty(cache.name)
self.check_calibrator_cleanup(calibrator)
def test_external_data(self):
with tempfile.NamedTemporaryFile(
) as path, tempfile.NamedTemporaryFile() as data:
model = OnnxFromPath(ONNX_MODELS["const_foldable"].path)
loader = SaveOnnx(model,
path.name,
external_data_path=data.name,
size_threshold=0)
loader()
check_file_non_empty(path.name)
check_file_non_empty(data.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])
check_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_pb(self):
with tempfile.NamedTemporaryFile() as outpath:
tf_loader = SaveGraph(GraphFromFrozen(TF_MODELS["identity"].path),
path=outpath.name)
tf_loader()
check_file_non_empty(outpath.name)
def test_tf2onnx_save_onnx(self):
with tempfile.NamedTemporaryFile() as outpath:
run_polygraphy_run([TF_MODELS["identity"].path, "--onnxrt", "--model-type=frozen", "--save-onnx", outpath.name])
check_file_non_empty(outpath.name)
import onnx
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:
check_file_non_empty(timeline)
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_tactic_replay(self):
with tempfile.NamedTemporaryFile() as tactic_replay:
run_polygraphy_run([ONNX_MODELS["identity"].path, "--trt", "--save-tactics", tactic_replay.name])
check_file_non_empty(tactic_replay.name)
run_polygraphy_run([ONNX_MODELS["identity"].path, "--trt", "--load-tactics", tactic_replay.name])
def test_save_load_engine(self):
with tempfile.NamedTemporaryFile() as outpath:
run_polygraphy_run([ONNX_MODELS["identity"].path, "--trt", "--save-engine", outpath.name])
check_file_non_empty(outpath.name)
run_polygraphy_run(["--trt", outpath.name, "--model-type=engine"])
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)
def test_save_onnx(self):
with tempfile.NamedTemporaryFile() as outpath:
loader = SaveOnnx(OnnxFromPath(ONNX_MODELS["identity"].path),
path=outpath.name)
loader()
check_file_non_empty(outpath.name)