def get_tensor_metadata(tensors):
metadata = TensorMetadata()
for tensor in tensors:
metadata.add(name=tensor.name,
dtype=get_dtype(tensor),
shape=get_shape(tensor))
return metadata
def get_input_metadata(self):
ONNX_RT_TYPE_TO_NP = {
"tensor(double)": np.float64,
"tensor(float)": np.float32,
"tensor(float16)": np.float16,
"tensor(int16)": np.int16,
"tensor(int32)": np.int32,
"tensor(int64)": np.int64,
"tensor(int8)": np.int8,
"tensor(uint16)": np.uint16,
"tensor(uint32)": np.uint32,
"tensor(uint64)": np.uint64,
"tensor(uint8)": np.uint8,
"tensor(bool)": np.bool,
}
def process_shape(shape):
# dims can be strings to indicate dynamic dimensions. Polygraphy uses None for the same purpose
return [
None if not isinstance(elem, int) else elem for elem in shape
]
meta = TensorMetadata()
for node in self.sess.get_inputs():
meta.add(node.name,
dtype=ONNX_RT_TYPE_TO_NP[node.type],
shape=process_shape(node.shape))
return meta
def get_layer_input_metadata(layer):
meta = TensorMetadata()
for i in range(layer.num_inputs):
inp = layer.get_input(i)
if inp:
meta.add(inp.name, np_dtype_from_trt(inp.dtype), inp.shape)
return meta
def get_layer_output_metadata(layer):
meta = TensorMetadata()
for i in range(layer.num_outputs):
outp = layer.get_output(i)
if outp:
meta.add(outp.name, np_dtype_from_trt(outp.dtype), outp.shape)
return meta
def get_input_metadata_from_profile(profile, network):
"""
Returns metadata about the inputs based on the OPT values set in a profile.
Args:
profile (trt.IOptimizationProfile):
The profile from which to retrieve input metada.
network (trt.INetworkDefinition):
The network the profile applies to.
Returns:
TensorMetadata:
A mapping of input names to their types and shapes.
Shapes are retrieved from the OPT values in the profile.
"""
input_metadata = TensorMetadata()
for index in range(network.num_inputs):
tensor = network.get_input(index)
if tensor.is_shape_tensor:
shapes = profile.get_shape_input(tensor.name)
else:
shapes = profile.get_shape(tensor.name)
if tuple(shapes[0]) != tuple(shapes[2]):
G_LOGGER.warning(
"Will use `opt` shapes from profile 0 for calibration. "
"Note that even though `min` != `max` in this profile, calibration "
"will use fixed input shapes (this is not necessarily an issue)."
)
# Always use opt shape
input_metadata.add(name=tensor.name,
dtype=trt.nptype(tensor.dtype),
shape=shapes[1])
return input_metadata
def _input_metadata_from_network(network):
input_metadata = TensorMetadata()
for index in range(network.num_inputs):
tensor = network.get_input(index)
input_metadata.add(name=tensor.name,
dtype=np.dtype(trt.nptype(tensor.dtype)),
shape=tensor.shape)
return input_metadata
def get_output_metadata(network):
outputs = TensorMetadata()
for i in range(network.num_outputs):
tensor = network.get_output(i)
outputs.add(name=tensor.name,
dtype=np_dtype_from_trt(tensor.dtype),
shape=tensor.shape)
return outputs
def metadata_from_names(names):
metadata = TensorMetadata()
for name in names:
dtype, shape = tensors.get(name, (None, None))
if name in initializer_metadata:
name = "Initializer | {:}".format(name)
metadata.add(name=name, dtype=dtype, shape=shape)
return metadata
def extract_model():
input_metadata = TensorMetadata().add("X",
dtype=np.float32,
shape=(64, 64))
output_metadata = TensorMetadata().add("identity_out_0",
dtype=np.float32,
shape=None)
return onnx_from_path(
ONNX_MODELS["identity_identity"].path), input_metadata, output_metadata
def get_tensor_metadata(tensors):
metadata = TensorMetadata()
for tensor in tensors:
try:
shape = [elem.value if hasattr(elem, "value") else elem for elem in tensor.shape]
except ValueError:
# Happens when rank is unknown
shape = None
metadata.add(tensor.name, dtype=tensor.dtype.as_numpy_dtype, shape=shape)
return metadata
def get_input_metadata_impl(self):
inputs = TensorMetadata()
for binding in self.engine:
if self.engine.binding_is_input(binding):
# Always prepend a dynamic batch dimension
inputs.add(binding,
trt.nptype(self.engine.get_binding_dtype(binding)),
[-1] + list(self.engine.get_binding_shape(binding)))
return inputs
def test_shape_tensor_detected(self):
INPUT_DATA = (1, 2, 3)
input_meta = TensorMetadata().add("X", dtype=np.int32, shape=(3, ))
# This contains the shape values
overriden_meta = TensorMetadata().add("X", dtype=np.int32, shape=INPUT_DATA)
data_loader = DataLoader(input_metadata=overriden_meta)
data_loader.input_metadata = input_meta
feed_dict = data_loader[0]
assert np.all(feed_dict["X"] == INPUT_DATA) # values become INPUT_DATA
def test_no_shape_tensor_false_positive_negative_dims(self):
INPUT_DATA = (-100, 2, 4)
# This should NOT be detected as a shape tensor
input_meta = TensorMetadata().add("X", dtype=np.int32, shape=(3, ))
overriden_meta = TensorMetadata().add("X", dtype=np.int32, shape=INPUT_DATA)
data_loader = DataLoader(input_metadata=overriden_meta)
data_loader.input_metadata = input_meta
feed_dict = data_loader[0]
assert feed_dict["X"].shape == (3, ) # Shape IS (3, ), because this is NOT a shape tensor
assert np.any(feed_dict["X"] != INPUT_DATA) # Contents are not INPUT_DATA, since it's not treated as a shape value
def test_no_shape_tensor_false_positive_float(self):
INPUT_DATA = (-100, -50, 0)
# Float cannot be a shape tensor
input_meta = TensorMetadata().add("X", dtype=np.float32, shape=(3, ))
overriden_meta = TensorMetadata().add("X", dtype=np.float32, shape=INPUT_DATA)
data_loader = DataLoader(input_metadata=overriden_meta)
data_loader.input_metadata = input_meta
feed_dict = data_loader[0]
assert feed_dict["X"].shape == (3, ) # Values are NOT (3, )
assert np.any(feed_dict["X"] != INPUT_DATA) # Values are NOT (3, )
def test_can_cast_dtype(self):
# Ensure that the data loader can only be used once
def load_data():
yield {"X": np.ones((1, 1), dtype=np.float32)}
cache = DataLoaderCache(load_data())
fp32_meta = TensorMetadata().add("X", dtype=np.float32, shape=(1, 1))
cache.set_input_metadata(fp32_meta)
feed_dict = cache[0]
assert feed_dict["X"].dtype == np.float32
fp64_meta = TensorMetadata().add("X", dtype=np.float64, shape=(1, 1))
cache.set_input_metadata(fp64_meta)
feed_dict = cache[0]
assert feed_dict["X"].dtype == np.float64
def fallback_shape_inference(onnx_model):
from polygraphy.backend.onnx import BytesFromOnnx, ModifyOnnx
from polygraphy.backend.onnxrt import (OnnxrtRunner,
SessionFromOnnxBytes)
load_model = ModifyOnnx(onnx_model, outputs=constants.MARK_ALL)
with OnnxrtRunner(SessionFromOnnxBytes(BytesFromOnnx(load_model))) as runner:
data_loader = self.makers[DataLoaderArgs].get_data_loader()
data_loader.input_metadata = runner.get_input_metadata()
outputs = runner.infer(feed_dict=data_loader[0])
meta = TensorMetadata()
for name, output in outputs.items():
meta.add(name, output.dtype, output.shape)
return meta
def test_range_min_max_equal(self):
RANGE_VAL = 1
data_loader = DataLoader(input_metadata=TensorMetadata().add(
"X", dtype=np.int32, shape=(1, 1)),
int_range=(RANGE_VAL, RANGE_VAL))
feed_dict = data_loader[0]
assert np.all(feed_dict["X"] == RANGE_VAL)
def fallback_inference(self, onnx_model):
"""
Run inference with ONNX-Runtime.
This can be used to retrieve values/shapes/data types for all
tensors in the model when other shape inference approaches fail.
Args:
onnx_model (onnx.ModelProto):
The ONNX model in which to infer shapes.
data_loader_args (DataLoaderArgs):
The data loader argument group to use to generate input data.
Returns:
(OrderedDict[str, np.ndarray], TensorMetadata):
1. Mapping of values for all tensors in the model, including inputs.
Values are loaded lazily when first accessed so as to save memory.
2. Metadata for every tensor in the model.
"""
from polygraphy.comparator import IterationResult
with G_LOGGER.verbosity(G_LOGGER.severity + 10):
load_model = onnx_backend.ModifyOutputs(onnx_model,
outputs=constants.MARK_ALL,
copy=True)
with onnxrt_backend.OnnxrtRunner(
onnxrt_backend.SessionFromOnnx(
onnx_backend.BytesFromOnnx(load_model))) as runner:
# We want to set input_metadata only - not user_input_metadata, so that user_input_metadata
# will be populated by the --model-inputs argument.
data_loader = self.data_loader_args.get_data_loader()
data_loader.input_metadata = runner.get_input_metadata()
feed_dict = data_loader[0]
with G_LOGGER.verbosity(G_LOGGER.severity - 10):
G_LOGGER.info(
"Running fallback shape inference using input metadata:\n{:}"
.format(TensorMetadata.from_feed_dict(feed_dict)))
outputs = runner.infer(feed_dict)
# We include the inputs here so that we have values for all tensors in the model.
outputs.update(feed_dict)
# Use IterationResult here since it can handle very large tensors by saving to disk.
# Layerwise outputs might otherwise take up too much memory.
return IterationResult(outputs), TensorMetadata.from_feed_dict(
outputs)
def test_non_user_provided_inputs_never_shape_tensors(self):
# If the user didn't provide metadata, then the value can never be a shape tensor.
input_meta = TensorMetadata().add("X", dtype=np.int32, shape=(3, ))
data_loader = DataLoader()
data_loader.input_metadata = input_meta
feed_dict = data_loader[0]
assert feed_dict["X"].shape == (3, ) # Treat as a normal tensor
def test_override_input_metadata(self):
arg_group = ArgGroupTestHelper(DataLoaderArgs(), deps=[ModelArgs()])
arg_group.parse_args([])
data_loader = arg_group.get_data_loader(
user_input_metadata=TensorMetadata().add(
"test0", dtype=np.float32, shape=(4, 4)))
for feed_dict in data_loader:
assert feed_dict["test0"].shape == (4, 4)
def parse(self, args):
def determine_model_type():
if args_util.get(args, "model_type") is not None:
return args.model_type.lower()
if args_util.get(args, "model_file") is None:
return None
def use_ext(ext_mapping):
file_ext = os.path.splitext(args.model_file)[-1]
if file_ext in ext_mapping:
return ext_mapping[file_ext]
runners = util.default(args_util.get(args, "runners"), [])
if args_util.get(args, "ckpt") or os.path.isdir(args.model_file):
return "ckpt"
elif "tf" in runners or "trt_legacy" in runners:
if args.caffe_model:
return "caffe"
return use_ext(ModelArgs.EXT_MODEL_TYPE_MAPPING) or "frozen"
else:
model_type = use_ext(ModelArgs.EXT_MODEL_TYPE_MAPPING)
if model_type:
return model_type
G_LOGGER.exit(
"Could not automatically determine model type for: {:}\n"
"Please explicitly specify the type with the --model-type option"
.format(args.model_file))
if args_util.get(args, "input_shapes"):
self.input_shapes = args_util.parse_meta(
args_util.get(args, "input_shapes"),
includes_dtype=False) # TensorMetadata
else:
self.input_shapes = TensorMetadata()
self.model_file = args_util.get(args, "model_file")
if self.model_file:
G_LOGGER.verbose("Model: {:}".format(self.model_file))
if not os.path.exists(self.model_file):
G_LOGGER.warning("Model path does not exist: {:}".format(
self.model_file))
self.model_file = os.path.abspath(self.model_file)
model_type_str = util.default(self._model_type, determine_model_type())
self.model_type = ModelArgs.ModelType(
model_type_str) if model_type_str else None
if self.model_type == "trt-network-script" and (
not self.model_file or not self.model_file.endswith(".py")):
G_LOGGER.exit(
"TensorRT network scripts must exist and have '.py' extensions. "
"Note: Provided network script path was: {:}".format(
self.model_file))
def test_can_override_shape(self):
model = ONNX_MODELS["dynamic_identity"]
shape = (1, 1, 4, 5)
custom_input_metadata = TensorMetadata().add("X", dtype=None, shape=shape)
data_loader = DataLoader(input_metadata=custom_input_metadata)
# Simulate what the comparator does
data_loader.input_metadata = model.input_metadata
feed_dict = data_loader[0]
assert tuple(feed_dict["X"].shape) == shape
def get_input_metadata_impl(self):
ONNX_RT_TYPE_TO_NP = {
"tensor(double)": np.float64,
"tensor(float)": np.float32,
"tensor(float16)": np.float16,
"tensor(int16)": np.int16,
"tensor(int32)": np.int32,
"tensor(int64)": np.int64,
"tensor(int8)": np.int8,
"tensor(uint16)": np.uint16,
"tensor(uint32)": np.uint32,
"tensor(uint64)": np.uint64,
"tensor(uint8)": np.uint8,
"tensor(bool)": np.bool,
"tensor(string)": np.unicode,
}
meta = TensorMetadata()
for node in self.sess.get_inputs():
dtype = ONNX_RT_TYPE_TO_NP[node.type] if node.type in ONNX_RT_TYPE_TO_NP else None
meta.add(node.name, dtype=dtype, shape=node.shape)
return meta
def missing_meta_tensors(input_metadata, output_metadata):
missing = TensorMetadata()
for name, (dtype, shape) in input_metadata.items():
if dtype is None or shape is None:
missing.add(name, dtype, shape)
for name, (dtype, shape) in output_metadata.items():
if dtype is None:
missing.add(name, dtype, shape)
return missing
def get_input_metadata_from_profile(profile, network):
"""
Returns metadata about the inputs based on a profile
Args:
profile (trt.IOptimizationProfile): The profile from which to retrieve input metada.
network (trt.INetworkDefinition): The network
Returns:
TensorMetadata: A mapping of input names to their types and shapes.
"""
input_metadata = TensorMetadata()
for index in range(network.num_inputs):
tensor = network.get_input(index)
if tensor.is_shape_tensor:
shapes = profile.get_shape_input(tensor.name)
else:
shapes = profile.get_shape(tensor.name)
if tuple(shapes[0]) != tuple(shapes[1]):
G_LOGGER.warning("In profile 0, min != max, using opt shapes for calibration")
# Always use opt shape
input_metadata.add(name=tensor.name, dtype=trt.nptype(tensor.dtype), shape=shapes[1])
return input_metadata
def parse_meta_new_impl(meta_args, includes_shape=True, includes_dtype=True):
SEP = ":"
meta = TensorMetadata()
for meta_arg in meta_args:
name, shape, dtype = None, None, None
def pop_meta(func):
nonlocal meta_arg
meta_arg, _, val = meta_arg.rpartition(SEP)
val = cast(val.strip())
if isinstance(val, str) and val.lower() == "auto":
return None
return func(val)
if includes_dtype:
dtype = pop_meta(func=np_type_from_str)
if includes_shape:
shape = pop_meta(func=lambda s: tuple(e for e in s if e != ""))
name = meta_arg
meta.add(name, dtype=dtype, shape=shape)
return meta
def execute_runner(runner, loader_cache):
with runner as active_runner:
input_metadata = active_runner.get_input_metadata()
G_LOGGER.info("{:35}\n---- Model Input(s) ----\n{:}".format(active_runner.name, input_metadata),
mode=LogMode.ONCE)
# DataLoaderCache will ensure that the feed_dict does not contain any extra entries
# based on the provided input_metadata.
loader_cache.set_input_metadata(input_metadata)
if warm_up:
G_LOGGER.start("{:35} | Running {:} warm-up run(s)".format(active_runner.name, warm_up))
try:
feed_dict = loader_cache[0]
except IndexError:
G_LOGGER.warning("{:} warm-up run(s) were requested, but data loader did not supply any data. "
"Skipping warm-up run(s)".format(warm_up))
else:
G_LOGGER.ultra_verbose("Warm-up Input Buffers:\n{:}".format(util.indent_block(feed_dict)))
# First do a few warm-up runs, and don't time them.
for _ in range(warm_up):
active_runner.infer(feed_dict=feed_dict)
G_LOGGER.finish("{:35} | Finished {:} warm-up run(s)".format(active_runner.name, warm_up))
# Then, actual iterations.
index = 0
iteration_results = []
total_runtime = 0
for index, feed_dict in enumerate(loader_cache):
G_LOGGER.extra_verbose(lambda: "{:35} | Feeding inputs:\n{:}".format(active_runner.name, util.indent_block(feed_dict)))
outputs = active_runner.infer(feed_dict=feed_dict)
runtime = active_runner.last_inference_time()
total_runtime += runtime
# Without a deep copy here, outputs will always reference the output of the last run
iteration_results.append(IterationResult(outputs=copy.deepcopy(outputs), runtime=runtime, runner_name=active_runner.name))
G_LOGGER.info(lambda: "{:35}\n---- Model Output(s) ----\n{:}".format(
active_runner.name, TensorMetadata().from_feed_dict(outputs)),
mode=LogMode.ONCE)
G_LOGGER.extra_verbose(lambda: "{:35} | Inference Time: {:.3f} ms | Received outputs:\n{:}".format(
active_runner.name, runtime * 1000.0, util.indent_block(outputs)))
total_runtime_ms = total_runtime * 1000.0
G_LOGGER.finish("{:35} | Completed {:} iteration(s) in {:.4g} ms | Average inference time: {:.4g} ms.".format(active_runner.name, index + 1, total_runtime_ms, total_runtime_ms / float(index + 1)))
return iteration_results
def test_will_not_give_up_on_first_cache_miss(self):
SHAPE = (32, 32)
DATA = [OrderedDict()]
DATA[0]["X"] = np.zeros(SHAPE, dtype=np.int64)
DATA[0]["Y"] = np.zeros(SHAPE, dtype=np.int64)
cache = DataLoaderCache(DATA)
cache.set_input_metadata(TensorMetadata().add("X", np.int64, shape=SHAPE).add("Y", np.int64, SHAPE))
# Populate the cache with bad X but good Y
cache.cache[0] = OrderedDict()
cache.cache[0]["X"] = np.ones((64, 64), dtype=np.int64)
cache.cache[0]["Y"] = np.ones(SHAPE, dtype=np.int64)
feed_dict = cache[0]
# Cache cannot reuse X, so it'll reload - we'll get all 0s from the data loader
assert np.all(feed_dict["X"] == 0)
# Cache can reuse Y, even though it's after X, so we'll get ones from the cache
assert np.all(feed_dict["Y"] == 1)
def meta_from_iter_result(iter_result):
meta = TensorMetadata()
for name, arr in iter_result.items():
meta.add(name, dtype=arr.dtype, shape=arr.shape)
return meta
def meta_from_gs_tensors(tensors):
"""Get TensorMetadata from a list of ONNX-GraphSurgeon tensors"""
meta = TensorMetadata()
for tensor in tensors:
meta.add(tensor.name, tensor.dtype, tensor.shape)
return meta
