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_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_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_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:
create_config = CreateConfig(int8=True, calibrator=Calibrator(data, cache=cache.name))
with EngineFromNetwork((builder, network), create_config)():
check_file_non_empty(cache.name)
def test_trt_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_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(misc.insert_suffix(outpath.name, "*"))
for timeline in timelines:
check_file_non_empty(timeline)
def test_calibrator_with_file_object_cache(self, identity_builder_network,
mode):
builder, network = identity_builder_network
data = [{"x": np.ones((1, 1, 2, 2), dtype=np.float32)}]
with tempfile.NamedTemporaryFile(mode=mode) as cache:
create_config = CreateConfig(int8=True,
calibrator=Calibrator(data,
cache=cache))
with func.invoke(
EngineFromNetwork((builder, network), create_config)):
if mode != "rb":
check_file_non_empty(cache.name)
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_save_engine(self, identity_network):
with tempfile.NamedTemporaryFile() as outpath:
engine_loader = SaveEngine(EngineFromNetwork(identity_network),
path=outpath.name)
with engine_loader() as engine:
check_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()
check_file_non_empty(outpath.name)
