def build_default_profile(builder, network, default_shape_value=None):
default_shape_value = misc.default_value(default_shape_value, DEFAULT_SHAPE_VALUE)
def override_shape(shape):
return tuple([default_shape_value if misc.is_dimension_dynamic(dim) else dim for dim in shape])
trt_profile = builder.create_optimization_profile()
for idx in range(network.num_inputs):
inp = network.get_input(idx)
with G_LOGGER.verbosity(G_LOGGER.CRITICAL): # WAR for spam from TRT
is_shape_tensor = inp.is_shape_tensor
if is_shape_tensor:
rank = inp.shape[0]
shape = (default_shape_value, ) * rank
G_LOGGER.warning("Input shape-tensor: {:24} | Will use input values: {:} in profile.\n"
"If this is incorrect, please provide a profile "
"that sets the values for this input shape-tensor.".format(inp.name, shape, rank), mode=LogMode.ONCE)
trt_profile.set_shape_input(inp.name, shape, shape, shape)
else:
shape = override_shape(inp.shape)
if override_shape(inp.shape) != inp.shape:
G_LOGGER.warning("Input tensor: {:24} | Will use shape: {:} in profile (tensor shape is: {:}).\n"
"If this is incorrect, please provide a profile "
"that sets the shape for this input tensor.".format(inp.name, shape, inp.shape), mode=LogMode.ONCE)
trt_profile.set_shape(inp.name, shape, shape, shape)
return check_profile(trt_profile)
def build_profile(builder, network, profile):
trt_profile = builder.create_optimization_profile()
unused_keys = set(profile.keys())
for idx in range(network.num_inputs):
inp = network.get_input(idx)
if inp.name in unused_keys:
unused_keys.remove(inp.name)
with G_LOGGER.verbosity(): # WAR for spam from TRT
is_shape_tensor = inp.is_shape_tensor
if is_shape_tensor:
if inp.name in profile:
shapes = profile[inp.name]
trt_profile.set_shape_input(inp.name, shapes.min, shapes.opt, shapes.max)
G_LOGGER.extra_verbose("Input shape-tensor: {:24} | Setting values to min: {:}, opt: {:}, max: {:}".format(inp.name, shapes.min, shapes.opt, shapes.max))
else:
G_LOGGER.warning("input shape-tensor: {:24} | No values provided. Assuming this is not a dynamic shape-tensor.".format(inp.name), mode=LogMode.ONCE)
elif misc.is_shape_dynamic(inp.shape):
shapes = profile[inp.name]
trt_profile.set_shape(inp.name, shapes.min, shapes.opt, shapes.max)
G_LOGGER.extra_verbose("Input tensor: {:24} | Setting shape to min: {:}, opt: {:}, max: {:}".format(inp.name, shapes.min, shapes.opt, shapes.max))
if unused_keys:
G_LOGGER.warning("Some inputs provided in the profile were unused: {:}".format(list(unused_keys)))
return check_profile(trt_profile)
def mark_layerwise(network):
# Layers within loops cannot be marked as network outputs.
LOOP_START_NAMES = ["TRIP_LIMIT", "ITERATOR", "RECURRENCE"]
LOOP_END_NAMES = ["LOOP_OUTPUT"]
LOOP_START_LAYERS = [getattr(trt.LayerType, attr) for attr in LOOP_START_NAMES if hasattr(trt.LayerType, attr)]
LOOP_END_LAYERS = [getattr(trt.LayerType, attr) for attr in LOOP_END_NAMES if hasattr(trt.LayerType, attr)]
EXCLUDE_OUTPUT_LAYERS = [trt.LayerType.SHAPE, trt.LayerType.CONSTANT]
outputs = []
in_loop = False
for layer in network:
if layer.type in LOOP_START_LAYERS:
G_LOGGER.warning("Loop detected. Please ensure the network is topologically sorted so that layers within "
"the loop body are not marked as network outputs in layerwise mode", mode=LogMode.ONCE)
in_loop = True
elif layer.type in LOOP_END_LAYERS:
in_loop = False
should_mark_layer = not in_loop and layer.type not in EXCLUDE_OUTPUT_LAYERS
if should_mark_layer:
for index in range(layer.num_outputs):
tensor = layer.get_output(index)
outputs.append(tensor.name)
G_LOGGER.verbose("Marking {:} tensors as outputs".format(len(outputs)))
mark_outputs(network, outputs)
def determine_format(shape):
"""
Guesses the data format of a given shape.
Args:
shape (Tuple[int]): The shape, including batch dimension.
Returns:
DataFormat: The determined data format.
"""
# The smaller this ratio, the closer a and b are.
def minmax_ratio(a, b):
return abs(max(a, b) / min(a, b))
# Assume all shapes include batch dimension
if len(shape) == 4:
# Typically, H and W are quite close, so if minmax_ratio(0, 1) > minmax_ratio(1, 2), then we assume CHW.
if minmax_ratio(shape[1], shape[2]) > minmax_ratio(shape[2], shape[3]):
return DataFormat.NCHW
return DataFormat.NHWC
elif len(shape) == 3:
return DataFormat.NHW
elif len(shape) == 2:
return DataFormat.NW
else:
G_LOGGER.warning("Cannot determine format for " + str(shape) +
". Currently only implemented for input_buffers with 1-3 non-batch dimensions. Please update this function!")
return DataFormat.UNKNOWN
def add_to_script(self, script):
script.add_import(imports=["TrtLegacyRunner"], frm="polygraphy.backend.trt_legacy")
G_LOGGER.warning("Legacy TensorRT runner only supports implicit batch TensorFlow/UFF, ONNX, and Caffe models")
if self.model_args.model_type == "onnx":
script.add_import(imports=["ParseNetworkFromOnnxLegacy"], frm="polygraphy.backend.trt_legacy")
onnx_loader = self.onnx_loader_args.add_onnx_loader(script, disable_outputs=True)
loader_name = script.add_loader(Script.format_str("ParseNetworkFromOnnxLegacy({:})", onnx_loader), "parse_network_from_onnx_legacy")
elif self.model_args.model_type == "caffe":
script.add_import(imports=["LoadNetworkFromCaffe"], frm="polygraphy.backend.trt_legacy")
loader_name = script.add_loader(Script.format_str("LoadNetworkFromCaffe({:}, {:}, {:}, {:})", self.model_args.model_file, self.caffe_model,
self.trt_outputs, self.batch_size), "parse_network_from_caffe")
else:
script.add_import(imports=["LoadNetworkFromUff"], frm="polygraphy.backend.trt_legacy")
if self.model_args.model_type == "uff":
script.add_import(imports=["LoadUffFile"], frm="polygraphy.backend.trt_legacy")
shapes = {name: shape for name, (_, shape) in self.trt_loader_args.input_shapes.items()}
loader_name = script.add_loader(Script.format_str("LoadUffFile({:}, {:}, {:})", self.model_args.model_file, misc.default_value(shapes, {}), self.trt_outputs), "load_uff_file")
else:
script.add_import(imports=["ConvertToUff"], frm="polygraphy.backend.trt_legacy")
loader_name = script.add_loader(Script.format_str("ConvertToUff({:}, save_uff={:}, preprocessor={:})", self.tf_loader_args.add_to_script(script), self.save_uff, self.preprocessor), "convert_to_uff")
loader_name = script.add_loader(Script.format_str("LoadNetworkFromUff({:}, uff_order={:})", loader_name, self.uff_order), "uff_network_loader")
runner_str = Script.format_str("TrtLegacyRunner({:}, {:}, {:}, fp16={:}, tf32={:}, load_engine={:}, save_engine={:}, layerwise={:}, plugins={:})",
loader_name, self.trt_loader_args.workspace, self.batch_size, self.trt_loader_args.fp16, self.trt_loader_args.tf32,
self.model_args.model_file if self.model_args.model_type == "engine" else None,
self.trt_runner_args.save_engine, self.trt_outputs==constants.MARK_ALL, self.trt_loader_args.plugins)
runner_name = script.add_loader(runner_str, "trt_legacy_runner")
script.add_runner(runner_name)
return runner_name
def __init__(self, sess, timeline_dir=None, name=None):
"""
Args:
sess (Callable() -> Tuple[tf.Session, Sequence[str]]):
A callable that can supply a tuple containing a
TensorFlow session and output names.
timeline_dir (str):
Path to write a TensorFlow timeline.
Note that profiling may affect execution time.
name (str):
The human-readable name prefix to use for this runner.
A runner count and timestamp will be appended to this prefix.
"""
super().__init__(name=name, prefix="tf-runner")
self._sess = sess
self.timeline_dir = timeline_dir
self.num_inferences = 0
self.run_options = None
self.run_metadata = None
if self.timeline_dir is not None:
# Enable profiling
G_LOGGER.warning(
"Profiling is enabled. This will impact performance")
self.run_options = tf.RunOptions(
trace_level=tf.RunOptions.FULL_TRACE)
self.run_metadata = tf.RunMetadata()
def attrs_to_dict(attrs):
attr_dict = OrderedDict()
for attr in attrs:
def process_attr(attr_str: str):
processed = getattr(attr, ONNX_PYTHON_ATTR_MAPPING[attr_str])
if attr_str == "STRING":
processed = processed.decode()
elif attr_str == "TENSOR":
tensor_str = "Tensor: [dtype={:}, shape={:}]".format(get_dtype(processed), get_shape(processed))
if mode == "full":
tensor_str += " | Values:\n" + misc.indent_block(str(get_values(processed)))
processed = tensor_str
elif attr_str == "GRAPH":
processed = "\n" + str_from_onnx_graph(processed, mode, tensors, indent_level=indent_level + 2)
elif attr_str == "FLOATS" or attr_str == "INTS":
# Proto hacky list to normal Python list
processed = [p for p in processed]
elif attr_str == "STRINGS":
processed = [p.decode() for p in processed]
return processed
if attr.type in ATTR_TYPE_MAPPING:
attr_str = ATTR_TYPE_MAPPING[attr.type]
if attr_str in ONNX_PYTHON_ATTR_MAPPING:
attr_dict[attr.name] = process_attr(attr_str)
else:
G_LOGGER.warning("Attribute of type {:} is currently unsupported. Skipping attribute.".format(attr_str))
else:
G_LOGGER.warning("Attribute type: {:} was not recognized. Was the graph generated with a newer IR "
"version than the installed `onnx` package? Skipping attribute.".format(attr.type))
return attr_dict
def write_calibration_cache(self, cache):
self.cache_contents = cache.tobytes()
self.has_cached_scales = True
if self._cache is None:
return
try:
if self._cache.seekable():
self._cache.seek(0)
bytes_written = self._cache.write(self.cache_contents)
if bytes_written != len(self.cache_contents):
G_LOGGER.warning(
"Could not write entire cache. Note: cache contains {:} bytes, but only "
"{:} bytes were written".format(
len(self.cache_contents), bytes_written))
except AttributeError:
G_LOGGER.info("Writing calibration cache to: {:}".format(
self._cache))
with open(self._cache, "wb") as f:
f.write(self.cache_contents)
except:
# Cache is not writable
return
else:
self._cache.flush()
def mark_layerwise(network):
# Layers within loops cannot be marked as network outputs.
LOOP_START_NAMES = ["TRIP_LIMIT", "ITERATOR"]
LOOP_END_NAMES = ["LOOP_OUTPUT"]
LOOP_START_LAYERS = [getattr(trt.LayerType, attr) for attr in LOOP_START_NAMES if hasattr(trt.LayerType, attr)]
LOOP_END_LAYERS = [getattr(trt.LayerType, attr) for attr in LOOP_END_NAMES if hasattr(trt.LayerType, attr)]
EXCLUDE_OUTPUT_LAYERS = [trt.LayerType.SHAPE, trt.LayerType.CONSTANT]
num_tensors_marked = 0
in_loop = False
for layer in network:
if layer.type in LOOP_START_LAYERS:
G_LOGGER.warning("Loop detected. Please ensure the network is topologically sorted so that layers within "
"the loop body are not marked as network outputs in layerwise mode")
in_loop = True
elif layer.type in LOOP_END_LAYERS:
in_loop = False
def should_mark_layer():
return not in_loop and layer.type not in EXCLUDE_OUTPUT_LAYERS
if should_mark_layer():
for index in range(layer.num_outputs):
tensor = layer.get_output(index)
if not tensor.is_network_output:
G_LOGGER.verbose("Marking {:} as an output".format(tensor.name))
network.mark_output(tensor)
num_tensors_marked += 1
G_LOGGER.verbose("Marking {:} tensors as outputs".format(num_tensors_marked))
def __init__(self, max_workspace_size=None, tf32=None, fp16=None, int8=None, profiles=None, calibrator=None, strict_types=None):
"""
Functor that creates a TensorRT IBuilderConfig.
Args:
max_workspace_size (int): The maximum workspace size, in bytes, when building the engine.
tf32 (bool): Whether to build the engine with TF32 precision enabled. Defaults to False.
fp16 (bool): Whether to build the engine with FP16 precision enabled. Defaults to False.
int8 (bool): Whether to build the engine with INT8 precision enabled. Defaults to False.
profiles (List[Profile]):
A list of optimization profiles to add to the configuration. Only needed for
networks with dynamic input shapes. If this is omitted for a network with
dynamic shapes, a default profile is created, where dynamic dimensions are
replaced with Polygraphy's DEFAULT_SHAPE_VALUE (defined in util/constants.py).
See `Profile` for details.
calibrator (trt.IInt8Calibrator):
An int8 calibrator. Only required in int8 mode when
the network does not have explicit precision. For networks with
dynamic shapes, the last profile provided (or default profile if
no profiles are provided) is used during calibration.
"""
self.max_workspace_size = misc.default_value(max_workspace_size, 1 << 24)
self.tf32 = misc.default_value(tf32, False)
self.fp16 = misc.default_value(fp16, False)
self.int8 = misc.default_value(int8, False)
self.profiles = misc.default_value(profiles, [])
self.calibrator = calibrator
self.strict_types = misc.default_value(strict_types, False)
if self.calibrator is not None and not self.int8:
G_LOGGER.warning("A calibrator was provided to `CreateConfig`, but int8 mode was not enabled. "
"Did you mean to set `int8=True` to enable building with int8 precision?")
def get_batch(self, names):
try:
host_buffers = next(self.data_loader_iter)
except StopIteration:
if not self.num_batches:
G_LOGGER.warning(
"Calibrator data loader provided no data. Possibilities include: (1) data loader "
"has no data to provide, (2) data loader was a generator, and the calibrator is being "
"reused across multiple loaders (generators cannot be rewound)"
)
return None
else:
self.num_batches += 1
for name, host_buffer in host_buffers.items():
if name not in self.device_buffers:
self.device_buffers[name] = DeviceBuffer(
shape=host_buffer.shape, dtype=host_buffer.dtype)
G_LOGGER.verbose("Allocated: {:}".format(
self.device_buffers[name]))
if self.num_batches > 1:
G_LOGGER.warning(
"The calibrator data loader provided an extra input ({:}) compared to the last set of inputs.\n"
"Should this input be removed, or did you accidentally omit an input before?"
.format(name))
device_buffer = self.device_buffers[name]
device_buffer.copy_from(host_buffer)
return [
device_buffer.address()
for device_buffer in self.device_buffers.values()
]
def read_calibration_cache(self):
def load_from_cache():
if self._cache is None:
return None
try:
if self._cache.seekable():
self._cache.seek(0)
return self._cache.read()
except AttributeError:
if os.path.exists(self._cache):
G_LOGGER.info(
"Reading calibration cache from: {:}".format(
self._cache),
mode=LogMode.ONCE)
with open(self._cache, "rb") as f:
return f.read()
except:
# Cache is not readable
return None
if not self.has_cached_scales:
self.cache_contents = load_from_cache()
if not self.cache_contents:
G_LOGGER.warning(
"Calibration cache was provided, but is empty. Will regenerate scales by running calibration.",
mode=LogMode.ONCE)
self.cache_contents = None
else:
self.has_cached_scales = True
return self.cache_contents
def log_mismatches(mismatches):
try:
with G_LOGGER.indent():
G_LOGGER.super_verbose("Mismatched indices:\n{:}".format(np.argwhere(mismatches)))
G_LOGGER.extra_verbose("Runner: {:40} | Mismatched values:\n{:}".format(iter_result0.runner_name, out0[mismatches]))
G_LOGGER.extra_verbose("Runner: {:40} | Mismatched values:\n{:}".format(iter_result1.runner_name, out1[mismatches]))
except:
G_LOGGER.warning("Failing to log mismatches - this may be because the outputs are of different shapes")
def __call__(self, builder, network):
"""
Creates a TensorRT IBuilderConfig that can be used by the EngineFromNetwork.
Args:
builder (trt.Builder):
The TensorRT builder to use to create the configuration.
network (trt.INetworkDefinition):
The TensorRT network for which to create the config. The network is used to
automatically create a default optimization profile if none are provided.
Returns:
trt.IBuilderConfig: The TensorRT builder configuration.
"""
with misc.FreeOnException([builder.create_builder_config()
]) as (config, ):
calibration_profile = None
for profile in self.profiles:
calibration_profile = trt_util.build_profile(
builder, network, profile)
config.add_optimization_profile(calibration_profile)
if not self.profiles:
calibration_profile = trt_util.build_default_profile(
builder, network)
config.add_optimization_profile(calibration_profile)
if self.profiles:
G_LOGGER.info("Configuring with profiles: {:}".format(
self.profiles))
config.max_workspace_size = int(self.max_workspace_size)
if self.strict_types:
config.set_flag(trt.BuilderFlag.STRICT_TYPES)
if not self.tf32:
with contextlib.suppress(AttributeError):
config.clear_flag(trt.BuilderFlag.TF32)
if self.fp16:
config.set_flag(trt.BuilderFlag.FP16)
if self.int8:
config.set_flag(trt.BuilderFlag.INT8)
if not network.has_explicit_precision:
if self.calibrator is not None:
input_metadata = trt_util.get_input_metadata_from_profile(
calibration_profile, network)
with contextlib.suppress(AttributeError):
self.calibrator.reset(input_metadata)
config.int8_calibrator = self.calibrator
with contextlib.suppress(AttributeError):
config.set_calibration_profile(calibration_profile)
else:
G_LOGGER.warning(
"Network does not have explicit precision and no calibrator was provided. Please ensure "
"that tensors in the network have dynamic ranges set, or provide a calibrator in order to use int8 mode."
)
return config
def mark_io(graph, attr, tensors, filter_const=True):
if filter_const:
tensors = [t for t in tensors if not isinstance(t, gs.Constant)]
if not tensors:
G_LOGGER.warning(
"No non-constant tensors are available to mark. "
"Try folding constants in the model with `polygraphy surgeon sanitize --fold-constants`"
)
setattr(graph, attr, tensors)
G_LOGGER.info("Marking model {attr}: {:}".format(getattr(graph, attr), attr=attr))
return graph
def infer_shapes(model):
try:
import onnx.shape_inference
except:
G_LOGGER.warning("Could not import onnx.shape_inference module, skipping shape inference")
try:
model = onnx.shape_inference.infer_shapes(model)
G_LOGGER.verbose("ONNX Shape Inference completed successfully")
except Exception as err:
G_LOGGER.warning("ONNX shape inference exited with an error:\n{:}".format(err))
finally:
return model
def check_model(model):
try:
import onnx
except:
G_LOGGER.warning("Could not import onnx module, skipping model check")
try:
onnx.checker.check_model(model)
G_LOGGER.verbose("ONNX Checker Passed")
except onnx.checker.ValidationError as err:
G_LOGGER.warning("ONNX Checker exited with an error:\n{:}".format(err))
finally:
return model
def __init__(self,
network_loader=None,
max_workspace_size=None,
max_batch_size=None,
fp16=None,
tf32=None,
load_engine=None,
save_engine=None,
layerwise=False,
plugins=[],
name=None):
"""
Creates a runner that manages a single TensorRT engine.
network_loader (BaseModelLoader):
A loader that returns a TRT builder, network, parser and input shapes.
max_workspace_size (int): The maximum workspace size.
max_batch_size (int): The maximum batch size.
fp16 (bool): Whether to run in fp16 mode
layerwise (bool): Whether to retrieve the outputs of every layer in the network.
name (str):
The human-readable name prefix to use for this runner.
A runner count and timestamp will be appended to this prefix.
"""
G_LOGGER.warning(
"TrtLegacyRunner is deprecated, and will be removed in a future release"
)
# Load any user-supplied plugin libraries. This must happen before everything else, including engine deserialization.
if plugins:
import ctypes
for plugin in plugins:
path = os.path.abspath(plugin)
G_LOGGER.info("Loading plugin library: {:}".format(path))
ctypes.CDLL(path)
# Choose a unique name for this runner.
super().__init__(name=name, prefix="trt-legacy-runner")
# Save parameters for activate and deactivate.
self.network_loader = network_loader
self.max_workspace_size = misc.default_value(max_workspace_size,
1 << 24)
self.fp16 = misc.default_value(fp16, False)
self.tf32 = misc.default_value(tf32, False)
self.load_engine = load_engine
self.engine_path = save_engine
self.layerwise = layerwise
self.max_batch_size = max_batch_size
def default_find_output_func(output_name, index, iter_result):
found_name = misc.find_in_dict(output_name, iter_result, index)
if found_name is None:
return None
elif found_name != output_name:
exact_match = misc.find_in_dict(found_name, iter_result0)
if exact_match == found_name:
G_LOGGER.verbose("Will not compare {:} with {:}, since the former already has an exact match: {:}".format(
found_name, output_name, exact_match))
return None # If the found output is being compared against another output already, skip this non-exact match
G_LOGGER.warning("Output names did not match exactly. Assuming {:} output: {:} "
"corresponds to output: {:}".format(
iter_result.runner_name, found_name, output_name))
return [found_name]
def deactivate(self):
"""
Deactivate the runner.
"""
if not self.is_active:
G_LOGGER.warning(
"Runner is not active; Will not deactivate. If you really want to "
"deactivate this runner, call deactivate_impl() directly")
return
try:
self.deactivate_impl()
finally:
self.is_active = False
def activate(self):
"""
Activate the runner for inference. This may involve allocating GPU buffers, for example.
"""
if self.is_active:
G_LOGGER.warning(
"Runner is already active; Will not activate again. If you really want to "
"activate this runner twice, call activate_impl() directly")
return
try:
self.activate_impl()
finally:
self.is_active = True
def get_static_shape(name, shape):
static_shape = shape
if misc.is_shape_dynamic(shape):
static_shape = misc.override_dynamic_shape(shape)
if static_shape != shape and name not in self.user_input_metadata:
if not misc.is_valid_shape_override(static_shape, shape):
G_LOGGER.critical(
"Input tensor: {:24} | Cannot override original shape: {:} to {:}"
.format(name, shape, static_shape))
G_LOGGER.warning(
"Input tensor: {:24} | Adjusted shape: {:} to: {:}. If this is incorrect, please set input_metadata "
"or provide a custom data loader.".format(
name, shape, static_shape),
mode=LogMode.ONCE)
return static_shape
def last_inference_time(self):
"""
Returns the total time required for the last call to ``infer()``.
Returns:
float: The time in seconds, or None if runtime was not measured by the runner.
"""
if self.inference_time is None:
G_LOGGER.warning(
"inference_time was not set by this runner. Inference time will be incorrect!"
"To correctly compare runtimes, please set the inference_time property in the"
"infer() function",
mode=LogMode.ONCE)
return None
return self.inference_time
def str_from_graph(graph, mode):
graph_str = ""
input_metadata = get_input_metadata(graph)
output_metadata = get_output_metadata(graph)
graph_str += "---- {:} Graph Inputs ----\n{:}\n\n".format(len(input_metadata), input_metadata)
graph_str += "---- {:} Graph Outputs ----\n{:}\n\n".format(len(output_metadata), output_metadata)
graph_str += "---- {:} Nodes ----\n".format(len(graph.as_graph_def().node))
if mode == "basic":
G_LOGGER.warning("Displaying layer information is unsupported for TensorFlow graphs. "
"Please use --mode=full if you would like to see the raw nodes")
if mode == "full":
for node in graph.as_graph_def().node:
graph_str += str(node) + "\n"
graph_str += "\n"
return misc.indent_block(graph_str, level=0)
def __init__(self, deploy, model, outputs, batch_size=None, dtype=None):
self.deploy = deploy
self.model = model
if not self.model:
G_LOGGER.warning(
"No model file provided for Caffe model, random weights will be used. To avoid this, "
"please set the model paramater, or --model")
if not outputs:
G_LOGGER.critical(
"Please set Caffe model outputs using the outputs parameter, or --trt-outputs. "
"Note: To determine possible outputs, try running: tail -n50 {:}"
.format(deploy))
self.outputs = outputs
self.dtype = misc.default_value(dtype, trt.float32)
self.batch_size = misc.default_value(batch_size, 1)
def set_input_metadata(self, input_metadata):
"""
Set the input metadata for the data loader.
Args:
input_metadata (TensorMetadata):
Input Metadata, including shape and type information. The cache may attempt to transform inputs to
match the specified input_metadata when data already in the cache does not exactly match.
"""
self.input_metadata = input_metadata
with contextlib.suppress(AttributeError):
self.data_loader.input_metadata = input_metadata
if not self.cache:
G_LOGGER.verbose("Loading inputs from data loader")
self.cache = list(self.data_loader)
if not self.cache:
G_LOGGER.warning("Data loader did not yield any input data.")
def get_shapes(lst, idx):
default_shapes = copy.copy(shapes)
if idx < len(lst):
default_shapes.update(parse_meta(lst[idx], includes_dtype=False))
# Don't care about dtype, and need to override dynamic dimensions
default_shapes = {
name: misc.override_dynamic_shape(shape)
for name, (_, shape) in default_shapes.items()
}
for name, (_, shape) in shapes.items():
if tuple(default_shapes[name]) != tuple(shape):
G_LOGGER.warning(
"Input tensor: {:} | For TensorRT profile, overriding shape: {:} to: {:}"
.format(name, shape, default_shapes[name]),
mode=LogMode.ONCE)
return default_shapes
def setup_parser(self, subparsers=None):
"""
Set up a command-line argument parser.
Args:
subparsers (argparse.SubParsers):
A subparser group from argparse, like that returned by ``ArgumentParser.add_subparsers()``.
If this is omitted, this function will generate a new ``ArgumentParser`` instance.
Defaults to None.
Returns:
argparse.ArgumentParser:
The newly created parser if ``subparsers`` is not provided, or the newly created subparser otherwise.
"""
assert self.__doc__, "No help output was provided for this tool!"
allow_abbrev = all(not maker.disable_abbrev for maker in self.arg_groups.values())
if subparsers is not None:
parser = subparsers.add_parser(
self.name, help=self.__doc__, add_help=True, description=self.__doc__, allow_abbrev=allow_abbrev
)
parser.set_defaults(subcommand=self)
else:
parser = argparse.ArgumentParser(add_help=True, description=self.__doc__, allow_abbrev=allow_abbrev)
for maker in self.arg_groups.values():
# Register each maker with every other maker
for other_maker in self.arg_groups.values():
maker.register(other_maker)
maker.check_registered()
# This must be done after registration, since some argument groups
# may conditionally define arguments based on what other groups are present.
for maker in self.arg_groups.values():
maker.add_to_parser(parser)
try:
self.add_parser_args(parser)
except Exception as err:
G_LOGGER.warning(
"Could not register tool argument parser for: {:}\nNote: Error was: {:}".format(self.name, err)
)
return parser
def __init__(self, graph, opset=None, optimize=None, fold_constant=None):
"""
Functor that loads a TensorFlow graph and converts it to ONNX using the tf2onnx converter.
Args:
graph (Callable() -> Tuple[tf.Graph, Sequence[str]]):
A callable that can supply a tuple containing a TensorFlow
graph and output names.
opset (int): The ONNX opset to use during conversion.
optimize (bool): Whether to use tf2onnx's graph optimization pass.
fold_constant (bool): Whether to fold constants in the TensorFlow Graph. Requires that ``optimize`` is also enabled. Defaults to True.
"""
self._graph = graph
self.opset = misc.default_value(opset, 11)
self.fold_constant = misc.default_value(fold_constant, True)
self.optimize = misc.default_value(optimize, True)
if self.fold_constant and not self.optimize:
G_LOGGER.warning("`fold_constant` is enabled, but `optimize` is disabled. Constant folding will not be performed")
def coerce_cached_input(index, name, dtype, shape):
cached_feed_dict = self.cache[iteration]
cached_name = misc.find_in_dict(name, cached_feed_dict, index)
assert cached_name is not None
if cached_name != name:
G_LOGGER.warning(
"Input tensor: {:24} | Cached buffer name ({:}) does not match input name ({:})."
.format(name, cached_name, name))
buffer = cached_feed_dict[cached_name]
if dtype != buffer.dtype:
G_LOGGER.warning(
"Input tensor: {:24} | Cached buffer dtype ({:}) does not match input dtype ({:}), attempting cast. "
.format(name, buffer.dtype,
np.dtype(dtype).name))
buffer = buffer.astype(dtype)
if not misc.is_valid_shape_override(buffer.shape, shape):
G_LOGGER.warning(
"Input tensor: {:24} | Cached buffer shape ({:}) does not match input shape ({:}), attempting reshape. "
.format(name, buffer.shape, shape))
buffer = misc.try_match_shape(buffer, shape)
assert buffer.dtype == dtype and misc.is_valid_shape_override(
buffer.shape, shape)
return buffer
