def test_tf_to_onnx_to_trt(self):
with tempfile.NamedTemporaryFile() as outmodel:
run_polygraphy_convert([
TF_MODELS["identity"].path, "--model-type=frozen",
"--convert-to=trt", "-o", outmodel.name
])
self.check_engine(outmodel.name)
def test_trt_network_config_script_to_engine(self):
script = dedent("""
from polygraphy.backend.trt import CreateNetwork, CreateConfig
from polygraphy import func
import tensorrt as trt
@func.extend(CreateNetwork())
def my_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)
@func.extend(CreateConfig())
def load_config(config):
config.set_flag(trt.BuilderFlag.FP16)
""")
with tempfile.NamedTemporaryFile(
"w+",
suffix=".py") as f, tempfile.NamedTemporaryFile() as outmodel:
f.write(script)
f.flush()
run_polygraphy_convert([
f.name, "--model-type=trt-network-script",
"--trt-network-func-name=my_load_network",
"--trt-config-script", f.name, "--convert-to=trt", "-o",
outmodel.name
])
self.check_engine(outmodel.name)
def test_onnx_to_trt(self):
with tempfile.NamedTemporaryFile(suffix=".engine") as outmodel:
run_polygraphy_convert([
ONNX_MODELS["identity"].path, "--model-type=onnx", "-o",
outmodel.name
])
self.check_engine(outmodel.name)
def test_tf2onnx(self):
with tempfile.NamedTemporaryFile(suffix=".onnx") as outmodel:
run_polygraphy_convert([
TF_MODELS["identity"].path, "--model-type=frozen", "-o",
outmodel.name
])
assert onnx.load(outmodel.name)
def test_fp_to_fp16(self):
with tempfile.NamedTemporaryFile() as outmodel:
run_polygraphy_convert([
ONNX_MODELS["identity_identity"].path, "--convert-to=onnx",
"--fp-to-fp16", "-o", outmodel.name
])
assert onnx.load(
outmodel.name
).graph.value_info[0].type.tensor_type.elem_type == 10
def test_modify_onnx_outputs(self):
with tempfile.NamedTemporaryFile(suffix=".onnx") as outmodel:
run_polygraphy_convert([
ONNX_MODELS["identity_identity"].path, "-o", outmodel.name,
"--onnx-outputs", "mark", "all"
])
model = onnx.load(outmodel.name)
assert len(model.graph.output) == 2
def test_onnx_to_trt_to_onnx_like(self, model_name):
with util.NamedTemporaryFile() as outmodel:
run_polygraphy_convert([
ONNX_MODELS[model_name].path,
"--convert-to=onnx-like-trt-network", "-o", outmodel.name
])
