def test_unmark_shape_outputs(self, reshape_network):
builder, network, parser = func.invoke(
ModifyNetwork(reshape_network,
outputs=constants.MARK_ALL,
exclude_outputs=["reduce_prod_out_gs_2"]))
with builder, network, parser:
assert network.num_outputs == 1
def test_loader_explicit_precision(self):
builder, network, parser = func.invoke(
NetworkFromOnnxBytes(ONNX_MODELS["identity"].loader,
explicit_precision=True))
with builder, network, parser:
assert not network.has_implicit_batch_dimension
assert network.has_explicit_precision
def test_mark_custom_outputs(self, identity_identity_network):
builder, network, parser = func.invoke(
ModifyNetwork(identity_identity_network,
outputs=["identity_out_0"]))
with builder, network, parser:
assert network.num_outputs == 1
assert network.get_output(0).name == "identity_out_0"
def test_mark_shape_outputs(self, reshape_network):
builder, network, parser = func.invoke(
ModifyNetwork(reshape_network,
outputs=["output", "reduce_prod_out_gs_2"]))
with builder, network, parser:
assert network.num_outputs == 2
assert network.get_output(0).name == "reduce_prod_out_gs_2"
assert network.get_output(0).is_shape_tensor
def test_load_graph(self):
with tf.compat.v1.Graph().as_default() as graph:
inp = tf.placeholder(shape=(1, 1, 1, 1), dtype=tf.float32)
out = tf.identity(inp)
graph, outputs = func.invoke(GraphFromFrozen(graph))
assert graph
assert outputs
def test_exclude_outputs_with_mark_layerwise(self,
identity_identity_network):
builder, network, parser = func.invoke(
ModifyNetwork(identity_identity_network,
outputs=constants.MARK_ALL,
exclude_outputs=["identity_out_2"]))
with builder, network, parser:
assert network.num_outputs == 1
assert network.get_output(0).name == "identity_out_0"
def __call__(self):
with misc.FreeOnException(
func.invoke(
CreateNetwork(
explicit_precision=self.explicit_precision,
explicit_batch=self.explicit_batch))) as (builder,
network):
parser = trt.OnnxParser(network, trt_util.TRT_LOGGER)
return builder, network, parser
def try_register_tool(module, tool_class):
global TOOL_REGISTRY
try:
toolmod = importlib.import_module(module)
tool = func.invoke(getattr(toolmod, tool_class))
TOOL_REGISTRY.append(tool)
except Exception as err:
TOOL_REGISTRY.append(MissingTool(tool_class.lower(), err=err))
def __call__(self):
"""
Builds a TensorRT engine.
Returns:
trt.ICudaEngine: The engine that was created.
"""
# If network is a callable, then we own its return value
ret, owns_network = misc.try_call(self._network)
builder, network, parser = misc.unpack_args(ret, num=3)
with contextlib.ExitStack() as stack:
if owns_network:
stack.enter_context(builder)
stack.enter_context(network)
if parser is not None:
stack.enter_context(parser)
else:
provided = "Builder and Network" if parser is None else "Builder, Network, and Parser"
G_LOGGER.verbose(
"{:} were provided directly instead of via a Callable. This loader will not assume ownership. "
"Please ensure that they are freed.".format(provided))
config, owns_config = misc.try_call(self._config, builder, network)
if owns_config:
stack.enter_context(config)
else:
G_LOGGER.verbose(
"Builder configuration was provided directly instead of via a Callable. This loader will not assume "
"ownership. Please ensure it is freed.")
network_log_mode = "full" if G_LOGGER.severity <= G_LOGGER.ULTRA_VERBOSE else "attrs"
G_LOGGER.super_verbose(
lambda: ("Displaying TensorRT Network:\n" + trt_util.
str_from_network(network, mode=network_log_mode)))
G_LOGGER.info("Building engine with configuration: {:}".format(
trt_util.str_from_config(config)))
if misc.version(trt.__version__) < misc.version("7.3"):
engine = builder.build_engine(network, config)
else:
engine = func.invoke(
EngineFromBytes(
builder.build_serialized_network(network, config)))
if hasattr(config.int8_calibrator, "free"):
# Must go before engine check to ensure calibrator is freed on failures too.
config.int8_calibrator.free()
if not engine:
G_LOGGER.critical(
"Invalid Engine. Please ensure the engine was built correctly"
)
return engine
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 func.invoke(
EngineFromNetwork((builder, network), create_config)):
check_file_non_empty(cache.name)
def run(self, args):
import tensorrt as trt
if not self.makers[TrtLoaderArgs].calibration_cache:
G_LOGGER.warning(
"Not using a calibration cache. Using a calibration cache may significantly speed up the search process"
)
self.precision = {
"float32": trt.float32,
"float16": trt.float16
}[args.precision]
if self.precision == trt.float16 and not self.makers[
TrtLoaderArgs].fp16:
self.makers[TrtLoaderArgs].fp16 = True
if self.precision == trt.float16 and not self.makers[
TrtLoaderArgs].int8:
G_LOGGER.warning(
"Using float16 as the higher precision, but float16 is also the lowest precision available. Did you mean to set --int8 as well?"
)
if not any([
self.makers[TrtLoaderArgs].tf32,
self.makers[TrtLoaderArgs].fp16,
self.makers[TrtLoaderArgs].int8
]):
G_LOGGER.critical(
"Please enable at least one precision besides float32 (e.g. --int8, --fp16)"
)
if self.makers[ModelArgs].model_type == "engine":
G_LOGGER.critical(
"The precision tool cannot work with engines, as they cannot be modified. "
"Please provide a different format, such as an ONNX or TensorFlow model."
)
self.args = args
self.golden = OrderedDict()
self.golden.update(misc.pickle_load(args.golden))
self.builder, self.network, self.parser = func.invoke(
self.makers[TrtLoaderArgs].get_trt_network_loader())
with self.builder, self.network, self.parser:
indices = self.find()
if indices is not None:
G_LOGGER.info(
"To achieve acceptable accuracy, try running layers: {:} in {:} precision"
.format(indices, self.precision))
else:
G_LOGGER.critical(
"Could not find a configuration that resulted in acceptable accuracy"
)
def inspect_trt(self, args):
from polygraphy.backend.trt import util as trt_util
if self.makers[ModelArgs].model_type == "engine":
if args.mode != "none":
G_LOGGER.warning(
"Displaying layer information for TensorRT engines is not currently supported"
)
with func.invoke(self.makers[TrtLoaderArgs].
get_trt_serialized_engine_loader()) as engine:
engine_str = trt_util.str_from_engine(engine)
G_LOGGER.info(
"==== TensorRT Engine ====\n{:}".format(engine_str))
else:
builder, network, parser = func.invoke(
self.makers[TrtLoaderArgs].get_trt_network_loader())
with builder, network, parser:
network_str = trt_util.str_from_network(
network, mode=args.mode).strip()
G_LOGGER.info(
"==== TensorRT Network ====\n{:}".format(network_str))
def test_calibrator_caches_without_explicit_cache(
self, identity_builder_network):
builder, network = identity_builder_network
data = [{"x": np.ones((1, 1, 2, 2), dtype=np.float32)}]
calibrator = Calibrator(data)
# First, populate the cache
create_config = CreateConfig(int8=True, calibrator=calibrator)
with func.invoke(EngineFromNetwork((builder, network), create_config)):
pass
# Check that the internal cache is populated
assert calibrator.read_calibration_cache()
def test_multithreaded_runners_from_engine(self):
model = ONNX_MODELS["identity"]
engine = func.invoke(
EngineFromNetwork(NetworkFromOnnxBytes(model.loader)))
with engine, TrtRunner(engine) as runner0, TrtRunner(
engine) as runner1:
t1 = threading.Thread(target=model.check_runner, args=(runner0, ))
t2 = threading.Thread(target=model.check_runner, args=(runner1, ))
t1.start()
t2.start()
t2.join()
t2.join()
def __call__(self):
"""
Freezes a TensorFlow graph, and folds constants.
Returns:
Tuple[tf.Graph, Sequence[str]]: The TensorFlow graph, and the names of its outputs.
"""
(graph, output_names), _ = misc.try_call(self._graph)
with tf.Session(graph=graph) as sess:
sess.run(tf.initializers.global_variables())
sess.run(tf.initializers.local_variables())
graphdef = sess.graph.as_graph_def()
removed = tf.graph_util.remove_training_nodes(graphdef)
G_LOGGER.ultra_verbose("Removed nodes: {:}".format(removed))
for node in graphdef.node:
if node.op == 'RefSwitch':
node.op = 'Switch'
for index in range(len(node.input)):
if 'moving_' in node.input[index]:
node.input[index] = node.input[index] + '/read'
elif node.op == 'AssignSub':
node.op = 'Sub'
if 'use_locking' in node.attr: del node.attr['use_locking']
elif node.op == 'AssignAdd':
node.op = 'Add'
if 'use_locking' in node.attr: del node.attr['use_locking']
elif node.op == 'Assign':
node.op = 'Identity'
if 'use_locking' in node.attr: del node.attr['use_locking']
if 'validate_shape' in node.attr:
del node.attr['validate_shape']
if len(node.input) == 2:
# input0: ref: Should be from a Variable node. May be uninitialized.
# input1: value: The value to be assigned to the variable.
node.input[0] = node.input[1]
del node.input[1]
# Strip port information from outputs
output_names = [name.split(":")[0] for name in output_names]
output_graph_def = tf.graph_util.convert_variables_to_constants(
sess, graphdef, output_names)
output_graph_def = self.constfold(output_graph_def, output_names)
return func.invoke(GraphFromFrozen(output_graph_def))
def test_calibrator_rechecks_cache_on_reset(self,
identity_builder_network):
builder, network = identity_builder_network
data = [{"x": np.ones((1, 1, 2, 2), dtype=np.float32)}]
with tempfile.NamedTemporaryFile(mode="wb+") as cache:
calibrator = Calibrator(data, cache=cache.name)
# First, populate the cache
create_config = CreateConfig(int8=True, calibrator=calibrator)
with func.invoke(
EngineFromNetwork((builder, network), create_config)):
pass
# Ensure that now the calibrator will read from the cache when reset
calibrator.reset()
assert not calibrator.has_cached_scales
assert len(calibrator.read_calibration_cache()) == os.stat(
cache.name).st_size
def test_calibrator_outside_polygraphy(self, identity_builder_network):
builder, network = identity_builder_network
NUM_BATCHES = 2
def generate_data():
for item in [np.ones(
(1, 1, 2, 2), dtype=np.float32)] * NUM_BATCHES:
yield {"x": item}
calibrator = Calibrator(generate_data())
config = builder.create_builder_config()
config.set_flag(trt.BuilderFlag.INT8)
config.int8_calibrator = calibrator
if misc.version(trt.__version__) < misc.version("7.3"):
engine = builder.build_engine(network, config)
else:
engine = func.invoke(
EngineFromBytes(
builder.build_serialized_network(network, config)))
with engine:
assert engine
def check_network(self, suffix):
"""
Checks whether the provided network is accurate compared to golden values.
Returns:
OrderedDict[str, OutputCompareResult]:
A mapping of output names to an object describing whether they matched, and what the
required tolerances were.
"""
from polygraphy.backend.trt import (EngineFromNetwork, ModifyNetwork,
SaveEngine, TrtRunner)
from polygraphy.comparator import Comparator, CompareFunc
with G_LOGGER.verbosity(severity=G_LOGGER.severity if self.args.
show_output else G_LOGGER.CRITICAL):
data_loader = self.makers[DataLoaderArgs].get_data_loader()
self.makers[
TrtLoaderArgs].strict_types = True # HACK: Override strict types so things actually run in the right precision.
config = func.invoke(
self.makers[TrtLoaderArgs].get_trt_config_loader(data_loader),
self.builder, self.network)
suffix = "-{:}-{:}".format(suffix, self.precision)
engine_path = misc.insert_suffix(
self.makers[TrtRunnerArgs].save_engine, suffix)
self.builder, self.network, self.parser = func.invoke(
ModifyNetwork((self.builder, self.network, self.parser),
outputs=self.makers[TrtLoaderArgs].outputs))
engine_loader = SaveEngine(EngineFromNetwork(
(self.builder, self.network, self.parser), config),
path=engine_path)
runners = [TrtRunner(engine_loader)]
results = Comparator.run(runners, data_loader=data_loader)
if self.makers[ComparatorCompareArgs].validate:
Comparator.validate(results)
results.update(self.golden)
compare_func = CompareFunc.basic_compare_func(
atol=self.makers[ComparatorCompareArgs].atol,
rtol=self.makers[ComparatorCompareArgs].rtol,
check_shapes=not self.makers[ComparatorCompareArgs].
no_shape_check)
accuracy_result = Comparator.compare_accuracy(
results, compare_func=compare_func)
tolerances = list(accuracy_result.values())[0][
0] # First iteration of first runner pair
for name, req_tol in tolerances.items():
if bool(req_tol):
G_LOGGER.finish(
"PASSED | Output: {:} | Required Tolerances: {:}".format(
name, req_tol))
else:
G_LOGGER.error(
"FAILED | Output: {:} | Required Tolerances: {:}".format(
name, req_tol))
return accuracy_result
def identity_builder_network():
builder, network, parser = func.invoke(
NetworkFromOnnxBytes(ONNX_MODELS["identity"].loader))
with builder, network, parser:
yield builder, network
def inspect_tf(self, args):
from polygraphy.backend.tf import util as tf_util
tf_graph, _ = func.invoke(self.makers[TfLoaderArgs].get_tf_loader())
graph_str = tf_util.str_from_graph(tf_graph, mode=args.mode).strip()
G_LOGGER.info("==== TensorFlow Graph ====\n{:}".format(graph_str))
def import_graph(self, args):
onnx_model = func.invoke(self.makers[OnnxLoaderArgs].get_onnx_loader())
return onnx_model, gs.import_onnx(onnx_model)
def inspect_onnx(self, args):
from polygraphy.backend.onnx import util as onnx_util
onnx_model = func.invoke(self.makers[OnnxLoaderArgs].get_onnx_loader())
model_str = onnx_util.str_from_onnx(onnx_model, mode=args.mode).strip()
G_LOGGER.info("==== ONNX Model ====\n{:}".format(model_str))
def test_shape_output(self):
model = ONNX_MODELS["reshape"]
engine = func.invoke(
EngineFromNetwork(NetworkFromOnnxBytes(model.loader)))
with engine, TrtRunner(engine.create_execution_context) as runner:
model.check_runner(runner)