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 util.NamedTemporaryFile(
"w+",
suffix=".py") as f, util.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 replay(request):
"""
Returns:
Tuple[FakeAlgorithmContext, Algorithm, FakeAlgorithm,
Union[str, TacticReplayData], Union[str, TacticReplayData]]:
This fixture returns 5 things:
1. A fake TensorRT algorithm context
2. A Polygraphy Algorithm instance
3. A fake TensorRT algorithm (with the same information as (2))
4. An input tactic replay data, populated with the Polygraphy Algorithm (2), either
as a ``TacticReplayData`` instance, or a path.
5. An output tactic replay data, empty, either as a ``TacticReplayData`` instance, or
a path.
"""
jsonify = request.param
name = "node_of_y"
context = fake_context(name)
trt_algo = fake_algo()
poly_algo = Algorithm.from_trt(context, trt_algo)
in_replay_data = TacticReplayData().add(name, poly_algo)
out_replay_data = TacticReplayData()
if jsonify:
inpath = util.NamedTemporaryFile("w")
in_replay_data.save(inpath.name)
in_replay_data = inpath.name
outpath = util.NamedTemporaryFile("r")
out_replay_data = outpath.name
yield context, poly_algo, trt_algo, in_replay_data, out_replay_data
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_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_config_script(self):
arg_group = ArgGroupTestHelper(TrtConfigArgs())
with util.NamedTemporaryFile("w+", suffix=".py") as f:
f.write(
dedent(
"""
from polygraphy.backend.trt import CreateConfig
from polygraphy import func
import tensorrt as trt
@func.extend(CreateConfig())
def my_load_config(config):
config.set_flag(trt.BuilderFlag.FP16)
"""
)
)
f.flush()
arg_group.parse_args(["--trt-config-script", f.name, "--trt-config-func-name=my_load_config"])
assert arg_group.trt_config_script == f.name
assert arg_group.trt_config_func_name == "my_load_config"
builder, network = create_network()
with builder, network, arg_group.create_config(builder, network) as config:
assert isinstance(config, trt.IBuilderConfig)
assert config.get_flag(trt.BuilderFlag.FP16)
def test_tf2onnx(self):
with util.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_import_from_script(self, loader):
script = dedent("""
from polygraphy.backend.trt import CreateNetwork
from polygraphy import func
import tensorrt as trt
@func.extend(CreateNetwork())
def 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)
""")
with util.NamedTemporaryFile("w+", suffix=".py") as f:
f.write(script)
f.flush()
if loader == InvokeFromScript:
load_network = loader(f.name, "load_network")
builder, network = load_network()
else:
builder, network = loader(f.name, "load_network")
with builder, network:
assert isinstance(builder, trt.Builder)
assert isinstance(network, trt.INetworkDefinition)
assert network.num_layers == 1
assert network.get_layer(0).type == trt.LayerType.IDENTITY
def test_data_loader_script(self):
arg_group = ArgGroupTestHelper(DataLoaderArgs())
with util.NamedTemporaryFile("w+", suffix=".py") as f:
f.write(
dedent(
"""
import numpy as np
def my_load_data():
for _ in range(5):
yield {"inp": np.ones((3, 5), dtype=np.float32) * 6.4341}
"""
)
)
f.flush()
arg_group.parse_args(["--data-loader-script", f.name, "--data-loader-func-name=my_load_data"])
assert arg_group.data_loader_script == f.name
assert arg_group.data_loader_func_name == "my_load_data"
data_loader = arg_group.get_data_loader()
data = list(data_loader)
assert len(data) == 5
assert all(np.all(d["inp"] == np.ones((3, 5), dtype=np.float32) * 6.4341) for d in data)
def test_tf_save_pb(self):
with util.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_onnx_to_trt(self):
with util.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_tf_to_onnx_to_trt(self):
with util.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_onnx_rt_save_onnx(self):
with util.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_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_serialized_engine_loader_from_lambda(self, identity_engine):
with util.NamedTemporaryFile() as outpath:
with open(outpath.name, "wb") as f, identity_engine.serialize() as buffer:
f.write(buffer)
loader = EngineFromBytes(lambda: open(outpath.name, "rb").read())
with loader() as engine:
assert isinstance(engine, trt.ICudaEngine)
def test_save_load_engine(self):
with util.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_no_opts(self):
with util.NamedTemporaryFile("w+", suffix=".py") as template:
run_polygraphy_template(["trt-config", "-o", template.name])
builder, network = create_network()
create_config = InvokeFromScript(template.name, "load_config")
with builder, network, create_config(builder, network) as config:
assert isinstance(config, trt.IBuilderConfig)
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_no_shape_inference_if_has_metadata(self):
with util.NamedTemporaryFile() as outmodel:
status = run_polygraphy_surgeon([
"extract", ONNX_MODELS["identity_identity"].path, "-o",
outmodel.name, "--inputs", "X:auto:auto"
])
onnx_model_sanity_check(outmodel.name)
assert not was_shape_inference_run(status)
def test_tf_save_timeline(self):
with util.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_modify_onnx_outputs(self):
with util.NamedTemporaryFile(suffix=".onnx") as outmodel:
run_polygraphy_surgeon([
"sanitize", ONNX_MODELS["identity_identity"].path, "-o",
outmodel.name, "--outputs", "mark", "all"
])
model = onnx.load(outmodel.name)
assert len(model.graph.output) == 2
def test_override_shapes_no_clear_const_tensors_meta(self):
with util.NamedTemporaryFile() as outmodel:
run_polygraphy_surgeon([
"sanitize",
ONNX_MODELS["const_foldable"].path,
"-o",
outmodel.name,
"--override-input-shapes=input:[1,3]",
])
def test_log_file(self):
logger = Logger()
with util.NamedTemporaryFile("w+") as log_file:
logger.log_file = log_file.name
assert logger.log_file == log_file.name
logger.info("Hello")
log_file.seek(0)
assert log_file.read() == "[I] Hello\n"
def test_tactic_replay(self, trt_config_args):
with util.NamedTemporaryFile(suffix=".json") as f:
trt_config_args.parse_args(["--tactic-replay", f.name])
builder, network = create_network()
with builder, network, trt_config_args.create_config(builder, network=network) as config:
recorder = config.algorithm_selector
assert recorder.make_func == TacticRecorder
assert recorder.path == f.name
def test_no_model_file(self):
with util.NamedTemporaryFile("w+", suffix=".py") as template:
run_polygraphy_template(["trt-network", "-o", template.name])
load_network = InvokeFromScript(template.name, "load_network")
builder, network = load_network()
with builder, network:
assert isinstance(builder, trt.Builder)
assert isinstance(network, trt.INetworkDefinition)
def test_fp_to_fp16(self):
with util.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_save_load_inputs(self):
with util.NamedTemporaryFile() as infile0, util.NamedTemporaryFile(
) as infile1:
run_polygraphy_run([
ONNX_MODELS["identity"].path, "--onnxrt", "--save-input-data",
infile0.name
])
run_polygraphy_run([
ONNX_MODELS["identity"].path,
"--onnxrt",
"--load-input-data",
infile0.name,
"--save-input-data",
infile1.name,
]) # Copy
run_polygraphy_run([
ONNX_MODELS["identity"].path, "--onnxrt", "--load-input-data",
infile0.name, infile1.name
])
def test_merge_inputs_outputs(self):
with util.NamedTemporaryFile() as inps, util.NamedTemporaryFile(
) as outs, util.NamedTemporaryFile() as merged:
run_polygraphy_run(
[
ONNX_MODELS["identity"].path, "--onnxrt", "--save-inputs",
inps.name, "--save-outputs", outs.name
],
disable_verbose=True,
)
run_polygraphy_data(
["to-input", inps.name, outs.name, "-o", merged.name])
merged_data = util.load_json(merged.name)
assert len(merged_data) == 1
assert list(merged_data[0].keys()) == ["x", "y"]
assert all(
isinstance(val, np.ndarray) for val in merged_data[0].values())
def test_polygraphy_run_gen_script(self):
with util.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_int8_calibration_cache(self):
with util.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_tactics(self, trt_config_args, opt, cls):
with util.NamedTemporaryFile("w+", suffix=".json") as f:
if opt == "--load-tactics":
TacticReplayData().save(f)
trt_config_args.parse_args([opt, f.name])
builder, network = create_network()
with builder, network, trt_config_args.create_config(builder, network=network) as config:
recorder = config.algorithm_selector
assert recorder.make_func == cls
assert recorder.path == f.name
