def deactivate(self):
"""
Deactivate the runner.
If the POLYGRAPHY_INTERNAL_CORRECTNESS_CHECKS environment variable is set to `1`, this
will also check that the runner was reset to its state prior to activation.
Generally, you should use a context manager instead of manually activating and deactivating.
For example:
::
with RunnerType(...) as runner:
runner.infer(...)
"""
if not self.is_active:
G_LOGGER.warning("{:35} | Not active; will not deactivate. If you really want to "
"deactivate this runner, call deactivate_impl() directly".format(self.name))
return
self.inference_time = None
self._cached_input_metadata = None
self.is_active = None
try:
self.deactivate_impl()
except:
raise # Needed so we can have the else clause
else:
self.is_active = False
if config.INTERNAL_CORRECTNESS_CHECKS:
old_state = self._pre_activate_runner_state
del self._pre_activate_runner_state
if old_state != vars(self):
G_LOGGER.internal_error("Runner state was not reset after deactivation. "
"Note:\nOld state: {:}\nNew state: {:}".format(old_state, vars(self)))
def deactivate(self):
"""
Deactivate the runner. For example, this may involve freeing CPU or GPU memory.
Generally, you should use a context manager instead of manually activating and deactivating.
For example:
::
with RunnerType(...) as runner:
runner.infer(...)
"""
if not self.is_active:
G_LOGGER.warning(
"{:35} | Not active; will not deactivate. If you really want to "
"deactivate this runner, call deactivate_impl() directly".
format(self.name))
return
self.inference_time = None
self.is_active = None
self.deactivate_impl()
self.is_active = False
if config.INTERNAL_CORRECTNESS_CHECKS:
old_state = self._pre_activate_runner_state
del self._pre_activate_runner_state
if old_state != vars(self):
G_LOGGER.internal_error(
"Runner state was not reset after deactivation. "
"Note:\nOld state: {:}\nNew state: {:}".format(
old_state, vars(self)))
def wrapped_func(node, intermediate_tensors):
inputs = []
for inp in node.inputs:
if inp.is_empty(): # Optional input
inputs.append(None)
elif isinstance(inp, gs.Constant):
inputs.append(inp.values)
elif inp.name in intermediate_tensors:
inputs.append(intermediate_tensors[inp.name])
else:
G_LOGGER.internal_error(
"Input: {:} was not found in intermediate tensors and is not a constant.\n"
"Note: Intermediate tensors include: {:}".format(
inp.name, list(intermediate_tensors.keys())))
outputs = func(node.attrs, *inputs)
if len(outputs) != len(node.outputs):
G_LOGGER.internal_error(
"{:} reference implementation returned the wrong number of outputs.\n"
"Note: Expected {:} but recevied {:}".format(
op, len(node.outputs), len(outputs)))
return {
out_tensor.name: out
for out_tensor, out in zip(node.outputs, outputs)
}
def __iadd__(self, other):
if not isinstance(other, Script.String):
G_LOGGER.internal_error(
"Cannot concatenate str and Script.String. Note: str was: {:}"
.format(other))
elif self.safe != other.safe:
G_LOGGER.internal_error(
"Cannot concatenate unsafe string ({:}) to safe string ({:})!"
.format(other, self.s))
self.s += other.s
return self
def ensure_safe(inp):
"""
Ensures that the input is marked as a safe string (i.e. Script.String(safe=True)).
"""
if config.INTERNAL_CORRECTNESS_CHECKS:
if not isinstance(inp, Script.String):
G_LOGGER.internal_error("Input to ensure_safe must be of type Script.String, but was: {:}".format(inp))
elif not inp.safe:
G_LOGGER.internal_error("Input string: {:} was not checked for safety. "
"This is a potential security risk!".format(inp))
return inp
def try_register_tool(module, tool_class):
global TOOL_REGISTRY
try:
toolmod = importlib.import_module(module)
ToolClass = getattr(toolmod, tool_class)
TOOL_REGISTRY.append(ToolClass())
except Exception as err:
G_LOGGER.internal_error(
"Could not load command-line tool: {:}.\nNote: Error was: {:}".
format(tool_class.lower(), err))
TOOL_REGISTRY.append(MissingTool(tool_class.lower(), err=err))
def wrapper(self, *args, **kwargs):
old_dict = copy.copy(vars(self))
ret = None
try:
ret = func(self, *args, **kwargs)
finally:
if vars(self) != old_dict:
G_LOGGER.internal_error(
"{:} was mutated in a constant method! Note:\nOld state: {:}\nNew state: {:}".format(
self, old_dict, vars(self)
)
)
return ret
def warn_deprecated(name, use_instead, remove_in, module_name=None):
if config.INTERNAL_CORRECTNESS_CHECKS and version(
polygraphy.__version__) >= version(remove_in):
G_LOGGER.internal_error(
"{:} should have been removed in version: {:}".format(
name, remove_in))
full_obj_name = "{:}.{:}".format(module_name,
name) if module_name else name
warnings.warn(
"{:} is deprecated and will be removed in Polygraphy {:}. "
"Use {:} instead.".format(full_obj_name, remove_in, use_instead),
DeprecationWarning,
stacklevel=3,
)
def __call__(self, pairs):
dct = OrderedDict(pairs)
if config.INTERNAL_CORRECTNESS_CHECKS:
custom_type_keys = [key for key in dct if key.startswith(TYPE_STRING_PREFIX)]
if custom_type_keys and custom_type_keys[0] not in self.polygraphy_registered:
G_LOGGER.internal_error("Custom type has no decode function registered! "
"Note: Encoded object is:\n{:}".format(dct))
# The encoder will insert special key-value pairs into dictionaries encoded from
# custom types. If we find one, then we know to decode using the corresponding custom
# type function.
for type_str, func in self.polygraphy_registered.items():
if type_str in dct and dct[type_str] == constants.TYPE_MARKER: # Found a custom type!
return func(dct)
return dct
def deprecate_impl(obj):
if config.INTERNAL_CORRECTNESS_CHECKS and version(
polygraphy.__version__) >= version(remove_in):
G_LOGGER.internal_error(
"{:} should have been removed in version: {:}".format(
obj, remove_in))
nonlocal name
name = name or obj.__name__
if inspect.ismodule(obj):
class DeprecatedModule(object):
def __getattr__(self, attr_name):
warn_deprecated(name, use_instead, remove_in, module_name)
self = obj
return getattr(self, attr_name)
def __setattr__(self, attr_name, value):
warn_deprecated(name, use_instead, remove_in, module_name)
self = obj
return setattr(self, attr_name, value)
DeprecatedModule.__doc__ = "Deprecated: Use {:} instead".format(
use_instead)
return DeprecatedModule()
elif inspect.isclass(obj):
class Deprecated(obj):
def __init__(self, *args, **kwargs):
warn_deprecated(name, use_instead, remove_in, module_name)
super().__init__(*args, **kwargs)
Deprecated.__doc__ = "Deprecated: Use {:} instead".format(
use_instead)
return Deprecated
elif inspect.isfunction(obj):
def wrapped(*args, **kwargs):
warn_deprecated(name, use_instead, remove_in, module_name)
return obj(*args, **kwargs)
wrapped.__doc__ = "Deprecated: Use {:} instead".format(use_instead)
return wrapped
else:
G_LOGGER.internal_error(
"deprecate is not implemented for: {:}".format(obj))
def _add_to_script(self, script, user_input_metadata_str=None):
needs_invoke = False
using_random_data = False
if self.data_loader_script:
script.add_import(imports=["mod"], frm="polygraphy")
data_loader = make_invocable("mod.import_from_script",
self.data_loader_script,
name=self.data_loader_func_name)
needs_invoke = True
elif self.load_inputs:
script.add_import(imports=["load_json"], frm="polygraphy.json")
data_loader = safe(
"[]\nfor input_data_path in {load_inputs}:"
"\n\t{data_loader}.extend(load_json(input_data_path, description='input data'))",
load_inputs=self.load_inputs,
data_loader=Script.DATA_LOADER_NAME,
)
else:
using_random_data = True
if user_input_metadata_str is None and self.model_args is not None and self.model_args.input_shapes:
user_input_metadata_str = self.model_args.input_shapes
if user_input_metadata_str:
script.add_import(imports=["TensorMetadata"],
frm="polygraphy.common")
data_loader = make_invocable_if_nondefault(
"DataLoader",
seed=self.seed,
iterations=self.iterations,
input_metadata=user_input_metadata_str,
int_range=self.int_range,
float_range=self.float_range,
val_range=self.val_range,
)
if data_loader:
script.add_import(imports=["DataLoader"],
frm="polygraphy.comparator")
if using_random_data != self.is_using_random_data():
G_LOGGER.internal_error(
"is_using_random_data() reported a false positive!")
return script.set_data_loader(data_loader), needs_invoke
def run(self, args):
G_LOGGER.start("Starting iterations")
builder, network, parser = util.unpack_args(
self.arg_groups[TrtNetworkLoaderArgs].load_network(), 3)
with contextlib.ExitStack() as stack:
stack.enter_context(builder)
stack.enter_context(network)
if parser:
stack.enter_context(parser)
self.setup(args, network)
num_passed = 0
num_total = 0
success = True
MAX_COUNT = 100000 # We don't want to loop forever. This many iterations ought to be enough for anybody.
for iteration in range(MAX_COUNT):
remaining = self.remaining()
G_LOGGER.start("RUNNING | Iteration {:}{:}".format(
iteration + 1,
" | Approximately {:} iteration(s) remaining".format(
remaining) if remaining is not None else "",
))
self.process_network(network, success)
try:
engine = self.arg_groups[TrtEngineLoaderArgs].build_engine(
(builder, network))
except Exception as err:
G_LOGGER.warning(
"Failed to create network or engine, continuing to the next iteration.\n"
"Note: Error was: {:}".format(err))
G_LOGGER.internal_error(
"Failed to create network or engine. See warning above for details."
)
success = False
else:
# Don't need to keep the engine around in memory - just serialize to disk and free it.
with engine:
self.arg_groups[TrtEngineSaveArgs].save_engine(
engine,
self.arg_groups[ArtifactSorterArgs].iter_artifact)
success = self.arg_groups[
ArtifactSorterArgs].sort_artifacts(iteration + 1)
num_total += 1
if success:
num_passed += 1
if self.stop(iteration, success):
break
else:
G_LOGGER.warning(
"Maximum number of iterations reached: {:}.\n"
"Iteration has been halted to prevent an infinite loop!".
format(MAX_COUNT))
G_LOGGER.finish(
"Finished {:} iteration(s) | Passed: {:}/{:} | Pass Rate: {:}%".
format(iteration + 1, num_passed, num_total,
float(num_passed) * 100 / float(num_total)))
def call_impl(self):
"""
Returns:
onnx.ModelProto: The new ONNX model with shapes inferred.
"""
model, _ = util.invoke_if_callable(self._model)
external_data_dir = self.external_data_dir
try:
if isinstance(model, onnx.ModelProto):
MODEL_SIZE = model.ByteSize()
if MODEL_SIZE > LARGE_MODEL_THRESHOLD:
G_LOGGER.warning(
"Attempting to run shape inference on a large model. "
"This may require a large amount of memory.\nIf memory consumption becomes too high, "
"the process may be killed. You may want to try disabling shape inference in that case. ",
mode=LogMode.ONCE,
)
if MODEL_SIZE > self.save_to_disk_threshold_bytes:
G_LOGGER.warning(
"Model size ({:.3} MiB) exceeds the in-memory size threshold: {:.3} MiB.\n"
"The model will be saved to a temporary file before shape inference is run."
.format(
MODEL_SIZE / (1024.0**2),
self.save_to_disk_threshold_bytes / (1024.0**2)),
mode=LogMode.ONCE,
)
outdir = tempfile.TemporaryDirectory()
outpath = os.path.join(outdir.name, "tmp_model.onnx")
save_onnx(model, outpath, external_data_path="ext.data")
model = outpath
external_data_dir = outdir.name
G_LOGGER.verbose("Starting ONNX shape inference")
if isinstance(model, onnx.ModelProto):
model = shape_inference.infer_shapes(model)
else:
tmp_path = util.NamedTemporaryFile(prefix="tmp_polygraphy_",
suffix=".onnx").name
G_LOGGER.verbose(
"Writing shape-inferred model to: {:}".format(tmp_path))
shape_inference.infer_shapes_path(model, tmp_path)
# When external_data_dir is unset, use the model's current directory
model = onnx_from_path(tmp_path,
external_data_dir=util.default(
external_data_dir,
os.path.dirname(model) or None))
G_LOGGER.verbose("ONNX Shape Inference completed successfully")
except Exception as err:
if not self.error_ok:
raise
G_LOGGER.warning(
"ONNX shape inference exited with an error:\n{:}".format(err))
G_LOGGER.internal_error(
"ONNX shape inference exited with an error:\n{:}".format(err))
if not isinstance(model, onnx.ModelProto):
model = onnx_from_path(
model, external_data_dir=self.external_data_dir)
return model
def __call__(self, args):
G_LOGGER.exit("Encountered an error when loading this tool:\n{:}".format(self.err))
def try_register_tool(module, tool_class):
global TOOL_REGISTRY
try:
toolmod = importlib.import_module(module)
ToolClass = getattr(toolmod, tool_class)
TOOL_REGISTRY.append(ToolClass())
except Exception as err:
TOOL_REGISTRY.append(MissingTool(tool_class.lower(), err=err))
try_register_tool("polygraphy.tools.run", "Run")
try_register_tool("polygraphy.tools.convert", "Convert")
try_register_tool("polygraphy.tools.inspect", "Inspect")
try_register_tool("polygraphy.tools.surgeon", "Surgeon")
try_register_tool("polygraphy.tools.template", "Template")
try_register_tool("polygraphy.tools.debug", "Debug")
try_register_tool("polygraphy.tools.to_json", "ToJSON")
# Check that tool names are unique
tool_names = [tool.name for tool in TOOL_REGISTRY]
duplicates = set([name for name in tool_names if tool_names.count(name) > 1])
if duplicates:
G_LOGGER.internal_error("Multiple tools have the same name. Duplicate tool names found: {:}".format(duplicates))
def call_impl(self):
"""
Returns:
onnx.ModelProto: The ONNX-like, but **not** valid ONNX, representation of the TensorRT network.
"""
ret, owns_network = util.invoke_if_callable(self._network)
builder, network, parser = util.unpack_args(ret, num=3)
if builder is None or network is None:
G_LOGGER.critical(
"Expected to recevie a (builder, network) tuple for the `network` parameter, "
"but received: ({:}, {:})".format(builder, network))
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)
tensor_map = {}
def tensors_from_meta(meta):
nonlocal tensor_map
tensors = []
for name, (dtype, shape) in meta.items():
if name not in tensor_map:
tensor_map[name] = gs.Variable(name=name,
dtype=dtype,
shape=shape)
tensors.append(tensor_map[name])
return tensors
nodes = []
graph_inputs = tensors_from_meta(
trt_util.get_network_input_metadata(network))
graph_outputs = tensors_from_meta(
trt_util.get_network_output_metadata(network))
LAYER_TYPE_CLASS_MAPPING = trt_util.get_layer_class_mapping()
for layer in network:
op_name = layer.type.name
if layer.type in LAYER_TYPE_CLASS_MAPPING:
layer.__class__ = LAYER_TYPE_CLASS_MAPPING[layer.type]
node_inputs = tensors_from_meta(
trt_util.get_layer_input_metadata(layer))
node_outputs = tensors_from_meta(
trt_util.get_layer_output_metadata(layer))
attrs = {}
attr_names = trt_util.get_layer_attribute_names(layer)
for name in attr_names:
with G_LOGGER.verbosity():
attr = getattr(layer, name)
if util.is_sequence(attr) or any(
isinstance(attr, cls)
for cls in [trt.Dims, trt.Permutation]):
try:
attr = list(attr)
except ValueError: # Invalid dims
attr = []
if hasattr(attr, "__entries"): # TensorRT Enums
attr = attr.name
if isinstance(attr, trt.ILoop):
attr = attr.name
VALID_TYPES = [np.ndarray, list, int, str, bool, float]
if not any(isinstance(attr, cls) for cls in VALID_TYPES):
G_LOGGER.internal_error(
"Unknown type: {:} for layer attribute: {:}.\n"
"Note: Layer was: {:}".format(
type(attr), attr, layer))
try:
attr = str(attr)
except:
attr = "<error during conversion>"
attrs[name] = attr
nodes.append(
gs.Node(name=layer.name,
op=op_name,
attrs=attrs,
inputs=node_inputs,
outputs=node_outputs))
graph = gs.Graph(name=network.name,
inputs=graph_inputs,
outputs=graph_outputs,
nodes=nodes)
return gs.export_onnx(graph)