Loaders

Module: polygraphy.backend.trt

class LoadPlugins(plugins=None, obj=None)[source]

Bases: polygraphy.backend.base.loader.BaseLoader

A passthrough loader that loads plugins from the specified paths. Passthrough here means that it can be used to wrap any other loader. The purpose of wrapping another loader is that you can control the order of execution when lazily evaluating.

For immediate evaluation, use load_plugins instead:

load_plugins(plugins=["/path/to/my/plugin.so", "/path/to/my/other_plugin.so"])

Loads plugins from the specified paths.

Parameters
  • plugins (List[str]) – A list of paths to plugin libraries to load before inference.

  • obj – An object or callable to return or call respectively. If obj is callable, extra parameters will be forwarded to obj. If obj is not callable, it will be returned.

call_impl(*args, **kwargs)[source]
Returns

The provided obj argument, or its return value if it is callable. Returns None if obj was not set.

Return type

object

__call__(*args, **kwargs)

Invokes the loader by forwarding arguments to call_impl.

Note: call_impl should not be called directly - use this function instead.

load_plugins(plugins=None, obj=None, *args, **kwargs)

Immediately evaluated functional variant of LoadPlugins .

Loads plugins from the specified paths.

Parameters
  • plugins (List[str]) – A list of paths to plugin libraries to load before inference.

  • obj – An object or callable to return or call respectively. If obj is callable, extra parameters will be forwarded to obj. If obj is not callable, it will be returned.

Returns

The provided obj argument, or its return value if it is callable. Returns None if obj was not set.

Return type

object

class CreateNetwork(explicit_batch=None)[source]

Bases: polygraphy.backend.base.loader.BaseLoader

Functor that creates an empty TensorRT network.

Creates an empty TensorRT network.

Parameters

explicit_batch (bool) – Whether to create the network with explicit batch mode. Defaults to True.

call_impl()[source]
Returns

The builder and empty network.

Return type

(trt.Builder, trt.INetworkDefinition)

__call__(*args, **kwargs)

Invokes the loader by forwarding arguments to call_impl.

Note: call_impl should not be called directly - use this function instead.

create_network(explicit_batch=None)

Immediately evaluated functional variant of CreateNetwork .

Creates an empty TensorRT network.

Parameters

explicit_batch (bool) – Whether to create the network with explicit batch mode. Defaults to True.

Returns

The builder and empty network.

Return type

(trt.Builder, trt.INetworkDefinition)

class NetworkFromOnnxBytes(model_bytes)[source]

Bases: polygraphy.backend.trt.loader.BaseNetworkFromOnnx

Functor that parses an ONNX model to create a trt.INetworkDefinition.

Parses an ONNX model.

Parameters

model_bytes (Union[bytes, Callable() -> bytes]) – A serialized ONNX model or a callable that returns one.

call_impl()[source]
Returns

A TensorRT network, as well as the builder used to create it, and the parser used to populate it.

Return type

(trt.IBuilder, trt.INetworkDefinition, trt.OnnxParser)

__call__(*args, **kwargs)

Invokes the loader by forwarding arguments to call_impl.

Note: call_impl should not be called directly - use this function instead.

network_from_onnx_bytes(model_bytes)

Immediately evaluated functional variant of NetworkFromOnnxBytes .

Parses an ONNX model.

Parameters

model_bytes (Union[bytes, Callable() -> bytes]) – A serialized ONNX model or a callable that returns one.

Returns

A TensorRT network, as well as the builder used to create it, and the parser used to populate it.

Return type

(trt.IBuilder, trt.INetworkDefinition, trt.OnnxParser)

class NetworkFromOnnxPath(path)[source]

Bases: polygraphy.backend.trt.loader.BaseNetworkFromOnnx

Functor that parses an ONNX model to create a trt.INetworkDefinition. This loader supports models with weights stored in an external location.

Parses an ONNX model from a file.

Parameters

path (str) – The path from which to load the model.

call_impl()[source]
Returns

A TensorRT network, as well as the builder used to create it, and the parser used to populate it.

Return type

(trt.IBuilder, trt.INetworkDefinition, trt.OnnxParser)

__call__(*args, **kwargs)

Invokes the loader by forwarding arguments to call_impl.

Note: call_impl should not be called directly - use this function instead.

network_from_onnx_path(path)

Immediately evaluated functional variant of NetworkFromOnnxPath .

Parses an ONNX model from a file.

Parameters

path (str) – The path from which to load the model.

Returns

A TensorRT network, as well as the builder used to create it, and the parser used to populate it.

Return type

(trt.IBuilder, trt.INetworkDefinition, trt.OnnxParser)

class PostprocessNetwork(network, func, name=None)[source]

Bases: polygraphy.backend.base.loader.BaseLoader

[EXPERIMENTAL] Functor that applies a given post-processing function to a TensorRT INetworkDefinition.

Applies a given post-processing function to a TensorRT INetworkDefinition.

Parameters
  • network (Union[Tuple[trt.Builder, trt.INetworkDefinition, Optional[parser]], Callable() -> Tuple[trt.Builder, trt.INetworkDefinition, Optional[parser]]) – A tuple containing a TensorRT builder, network and optionally parser or a callable that returns one. To omit the parser, return a tuple containing just the builder and network.

  • func (Callable[[trt.INetworkDefinition], None]) – A callable which accepts a named network argument. PostprocessNetwork will pass in the parsed network via this argument, which can then be modified by the callable.

  • name (Optional[str]) – The name of this postprocessing step, used for logging purposes.

call_impl()[source]
Returns

The modified network along with the builder and parser if provided.

Return type

Tuple[trt.Builder, trt.INetworkDefinition, Optional[parser]

__call__(*args, **kwargs)

Invokes the loader by forwarding arguments to call_impl.

Note: call_impl should not be called directly - use this function instead.

postprocess_network(network, func, name=None)

Immediately evaluated functional variant of PostprocessNetwork .

Applies a given post-processing function to a TensorRT INetworkDefinition.

Parameters
  • network (Union[Tuple[trt.Builder, trt.INetworkDefinition, Optional[parser]], Callable() -> Tuple[trt.Builder, trt.INetworkDefinition, Optional[parser]]) – A tuple containing a TensorRT builder, network and optionally parser or a callable that returns one. To omit the parser, return a tuple containing just the builder and network.

  • func (Callable[[trt.INetworkDefinition], None]) – A callable which accepts a named network argument. PostprocessNetwork will pass in the parsed network via this argument, which can then be modified by the callable.

  • name (Optional[str]) – The name of this postprocessing step, used for logging purposes.

Returns

The modified network along with the builder and parser if provided.

Return type

Tuple[trt.Builder, trt.INetworkDefinition, Optional[parser]

class ModifyNetworkOutputs(network, outputs=None, exclude_outputs=None)[source]

Bases: polygraphy.backend.trt.loader.PostprocessNetwork

Functor that modifies outputs in a TensorRT INetworkDefinition.

Modifies outputs in a TensorRT INetworkDefinition.

Parameters
  • network (Union[Tuple[trt.Builder, trt.INetworkDefinition, Optional[parser]], Callable() -> Tuple[trt.Builder, trt.INetworkDefinition, Optional[parser]]) – A tuple containing a TensorRT builder, network and optionally parser or a callable that returns one. To omit the parser, return a tuple containing just the builder and network.

  • outputs (Sequence[str]) – Names of tensors to mark as outputs. If provided, this will override the outputs already marked in the network. If a value of constants.MARK_ALL is used instead of a list, all tensors in the network are marked.

  • exclude_outputs (Sequence[str]) – Names of tensors to exclude as outputs. This can be useful in conjunction with outputs=constants.MARK_ALL to omit outputs.

__call__(*args, **kwargs)

Invokes the loader by forwarding arguments to call_impl.

Note: call_impl should not be called directly - use this function instead.

call_impl()
Returns

The modified network along with the builder and parser if provided.

Return type

Tuple[trt.Builder, trt.INetworkDefinition, Optional[parser]

modify_network_outputs(network, outputs=None, exclude_outputs=None)

Immediately evaluated functional variant of ModifyNetworkOutputs .

Modifies outputs in a TensorRT INetworkDefinition.

Parameters
  • network (Union[Tuple[trt.Builder, trt.INetworkDefinition, Optional[parser]], Callable() -> Tuple[trt.Builder, trt.INetworkDefinition, Optional[parser]]) – A tuple containing a TensorRT builder, network and optionally parser or a callable that returns one. To omit the parser, return a tuple containing just the builder and network.

  • outputs (Sequence[str]) – Names of tensors to mark as outputs. If provided, this will override the outputs already marked in the network. If a value of constants.MARK_ALL is used instead of a list, all tensors in the network are marked.

  • exclude_outputs (Sequence[str]) – Names of tensors to exclude as outputs. This can be useful in conjunction with outputs=constants.MARK_ALL to omit outputs.

Returns

The modified network along with the builder and parser if provided.

Return type

Tuple[trt.Builder, trt.INetworkDefinition, Optional[parser]

class SetLayerPrecisions(network, layer_precisions)[source]

Bases: polygraphy.backend.trt.loader.PostprocessNetwork

Functor that sets layer precisions in a TensorRT INetworkDefinition.

Sets layer precisions in a TensorRT INetworkDefinition.

Parameters
  • network (Union[Tuple[trt.Builder, trt.INetworkDefinition, Optional[parser]], Callable() -> Tuple[trt.Builder, trt.INetworkDefinition, Optional[parser]]) – A tuple containing a TensorRT builder, network and optionally parser or a callable that returns one. To omit the parser, return a tuple containing just the builder and network.

  • layer_precisions (Dict[str, trt.DataType]) – A mapping of layer names to their desired compute precision.

__call__(*args, **kwargs)

Invokes the loader by forwarding arguments to call_impl.

Note: call_impl should not be called directly - use this function instead.

call_impl()
Returns

The modified network along with the builder and parser if provided.

Return type

Tuple[trt.Builder, trt.INetworkDefinition, Optional[parser]

set_layer_precisions(network, layer_precisions)

Immediately evaluated functional variant of SetLayerPrecisions .

Sets layer precisions in a TensorRT INetworkDefinition.

Parameters
  • network (Union[Tuple[trt.Builder, trt.INetworkDefinition, Optional[parser]], Callable() -> Tuple[trt.Builder, trt.INetworkDefinition, Optional[parser]]) – A tuple containing a TensorRT builder, network and optionally parser or a callable that returns one. To omit the parser, return a tuple containing just the builder and network.

  • layer_precisions (Dict[str, trt.DataType]) – A mapping of layer names to their desired compute precision.

Returns

The modified network along with the builder and parser if provided.

Return type

Tuple[trt.Builder, trt.INetworkDefinition, Optional[parser]

class SetTensorDatatypes(network, tensor_datatypes)[source]

Bases: polygraphy.backend.trt.loader.PostprocessNetwork

Functor that sets tensor datatypes in a TensorRT INetworkDefinition.

Sets tensor datatypes in a TensorRT INetworkDefinition.

Parameters
  • network (Union[Tuple[trt.Builder, trt.INetworkDefinition, Optional[parser]], Callable() -> Tuple[trt.Builder, trt.INetworkDefinition, Optional[parser]]) – A tuple containing a TensorRT builder, network and optionally parser or a callable that returns one. To omit the parser, return a tuple containing just the builder and network.

  • tensor_datatypes (Dict[str, trt.DataType]) – A mapping of tensor names to their desired data types.

__call__(*args, **kwargs)

Invokes the loader by forwarding arguments to call_impl.

Note: call_impl should not be called directly - use this function instead.

call_impl()
Returns

The modified network along with the builder and parser if provided.

Return type

Tuple[trt.Builder, trt.INetworkDefinition, Optional[parser]

set_tensor_datatypes(network, tensor_datatypes)

Immediately evaluated functional variant of SetTensorDatatypes .

Sets tensor datatypes in a TensorRT INetworkDefinition.

Parameters
  • network (Union[Tuple[trt.Builder, trt.INetworkDefinition, Optional[parser]], Callable() -> Tuple[trt.Builder, trt.INetworkDefinition, Optional[parser]]) – A tuple containing a TensorRT builder, network and optionally parser or a callable that returns one. To omit the parser, return a tuple containing just the builder and network.

  • tensor_datatypes (Dict[str, trt.DataType]) – A mapping of tensor names to their desired data types.

Returns

The modified network along with the builder and parser if provided.

Return type

Tuple[trt.Builder, trt.INetworkDefinition, Optional[parser]

class SetTensorFormats(network, tensor_formats)[source]

Bases: polygraphy.backend.trt.loader.PostprocessNetwork

Functor that sets tensor formats in a TensorRT INetworkDefinition.

Sets tensor formats in a TensorRT INetworkDefinition.

Parameters
  • network (Union[Tuple[trt.Builder, trt.INetworkDefinition, Optional[parser]], Callable() -> Tuple[trt.Builder, trt.INetworkDefinition, Optional[parser]]) – A tuple containing a TensorRT builder, network and optionally parser or a callable that returns one. To omit the parser, return a tuple containing just the builder and network.

  • tensor_formats (Dict[str, List[trt.TensorFormat]]) – A mapping of tensor names to their allowed formats.

__call__(*args, **kwargs)

Invokes the loader by forwarding arguments to call_impl.

Note: call_impl should not be called directly - use this function instead.

call_impl()
Returns

The modified network along with the builder and parser if provided.

Return type

Tuple[trt.Builder, trt.INetworkDefinition, Optional[parser]

set_tensor_formats(network, tensor_formats)

Immediately evaluated functional variant of SetTensorFormats .

Sets tensor formats in a TensorRT INetworkDefinition.

Parameters
  • network (Union[Tuple[trt.Builder, trt.INetworkDefinition, Optional[parser]], Callable() -> Tuple[trt.Builder, trt.INetworkDefinition, Optional[parser]]) – A tuple containing a TensorRT builder, network and optionally parser or a callable that returns one. To omit the parser, return a tuple containing just the builder and network.

  • tensor_formats (Dict[str, List[trt.TensorFormat]]) – A mapping of tensor names to their allowed formats.

Returns

The modified network along with the builder and parser if provided.

Return type

Tuple[trt.Builder, trt.INetworkDefinition, Optional[parser]

class EngineBytesFromNetwork(network, config=None, save_timing_cache=None)[source]

Bases: polygraphy.backend.base.loader.BaseLoader

Functor that uses a TensorRT INetworkDefinition to build a serialized engine.

Builds and serializes TensorRT engine.

Parameters
  • network (Union[Tuple[trt.Builder, trt.INetworkDefinition, Optional[parser]], Callable() -> Tuple[trt.Builder, trt.INetworkDefinition, Optional[parser]]) – A tuple containing a TensorRT builder, network and optionally parser or a callable that returns one. To omit the parser, return a tuple containing just the builder and network.

  • config (Callable(trt.Builder, trt.INetworkDefinition) -> trt.IBuilderConfig) – A TensorRT builder configuration or a callable that returns one. If not supplied, a CreateConfig instance with default parameters is used.

  • save_timing_cache (Union[str, file-like]) – A path or file-like object at which to save a tactic timing cache. Any existing cache will be appended to. If a path is provided, the file will be locked for exclusive access to prevent multiple processes from attempting to update the timing cache at the same time.

call_impl()[source]
Returns

The serialized engine that was created.

Return type

bytes

__call__(*args, **kwargs)

Invokes the loader by forwarding arguments to call_impl.

Note: call_impl should not be called directly - use this function instead.

engine_bytes_from_network(network, config=None, save_timing_cache=None)

Immediately evaluated functional variant of EngineBytesFromNetwork .

Builds and serializes TensorRT engine.

Parameters
  • network (Union[Tuple[trt.Builder, trt.INetworkDefinition, Optional[parser]], Callable() -> Tuple[trt.Builder, trt.INetworkDefinition, Optional[parser]]) – A tuple containing a TensorRT builder, network and optionally parser or a callable that returns one. To omit the parser, return a tuple containing just the builder and network.

  • config (Callable(trt.Builder, trt.INetworkDefinition) -> trt.IBuilderConfig) – A TensorRT builder configuration or a callable that returns one. If not supplied, a CreateConfig instance with default parameters is used.

  • save_timing_cache (Union[str, file-like]) – A path or file-like object at which to save a tactic timing cache. Any existing cache will be appended to. If a path is provided, the file will be locked for exclusive access to prevent multiple processes from attempting to update the timing cache at the same time.

Returns

The serialized engine that was created.

Return type

bytes

class EngineFromNetwork(network, config=None, save_timing_cache=None)[source]

Bases: polygraphy.backend.trt.loader.EngineBytesFromNetwork

Similar to EngineBytesFromNetwork, but returns an ICudaEngine instance instead of a serialized engine.

Builds and serializes TensorRT engine.

Parameters
  • network (Union[Tuple[trt.Builder, trt.INetworkDefinition, Optional[parser]], Callable() -> Tuple[trt.Builder, trt.INetworkDefinition, Optional[parser]]) – A tuple containing a TensorRT builder, network and optionally parser or a callable that returns one. To omit the parser, return a tuple containing just the builder and network.

  • config (Callable(trt.Builder, trt.INetworkDefinition) -> trt.IBuilderConfig) – A TensorRT builder configuration or a callable that returns one. If not supplied, a CreateConfig instance with default parameters is used.

  • save_timing_cache (Union[str, file-like]) – A path or file-like object at which to save a tactic timing cache. Any existing cache will be appended to. If a path is provided, the file will be locked for exclusive access to prevent multiple processes from attempting to update the timing cache at the same time.

call_impl()[source]
Returns

The engine that was created.

Return type

trt.ICudaEngine

__call__(*args, **kwargs)

Invokes the loader by forwarding arguments to call_impl.

Note: call_impl should not be called directly - use this function instead.

engine_from_network(network, config=None, save_timing_cache=None)

Immediately evaluated functional variant of EngineFromNetwork .

Builds and serializes TensorRT engine.

Parameters
  • network (Union[Tuple[trt.Builder, trt.INetworkDefinition, Optional[parser]], Callable() -> Tuple[trt.Builder, trt.INetworkDefinition, Optional[parser]]) – A tuple containing a TensorRT builder, network and optionally parser or a callable that returns one. To omit the parser, return a tuple containing just the builder and network.

  • config (Callable(trt.Builder, trt.INetworkDefinition) -> trt.IBuilderConfig) – A TensorRT builder configuration or a callable that returns one. If not supplied, a CreateConfig instance with default parameters is used.

  • save_timing_cache (Union[str, file-like]) – A path or file-like object at which to save a tactic timing cache. Any existing cache will be appended to. If a path is provided, the file will be locked for exclusive access to prevent multiple processes from attempting to update the timing cache at the same time.

Returns

The engine that was created.

Return type

trt.ICudaEngine

class EngineFromBytes(serialized_engine)[source]

Bases: polygraphy.backend.base.loader.BaseLoader

Functor that deserializes an engine from a buffer.

Deserializes an engine from a buffer.

Parameters

serialized_engine (Union[Union[str, bytes], Callable() -> Union[str, bytes]]) – The serialized engine bytes or a callable that returns them.

call_impl()[source]
Returns

The deserialized engine.

Return type

trt.ICudaEngine

__call__(*args, **kwargs)

Invokes the loader by forwarding arguments to call_impl.

Note: call_impl should not be called directly - use this function instead.

engine_from_bytes(serialized_engine)

Immediately evaluated functional variant of EngineFromBytes .

Deserializes an engine from a buffer.

Parameters

serialized_engine (Union[Union[str, bytes], Callable() -> Union[str, bytes]]) – The serialized engine bytes or a callable that returns them.

Returns

The deserialized engine.

Return type

trt.ICudaEngine

class BytesFromEngine(engine)[source]

Bases: polygraphy.backend.base.loader.BaseLoader

Functor that serializes an engine.

Serializes an engine.

Parameters

engine (Union[trt.ICudaEngine, Callable() -> trt.ICudaEngine]) – An engine or a callable that returns one.

call_impl()[source]
Returns

The serialized engine.

Return type

bytes

__call__(*args, **kwargs)

Invokes the loader by forwarding arguments to call_impl.

Note: call_impl should not be called directly - use this function instead.

bytes_from_engine(engine)

Immediately evaluated functional variant of BytesFromEngine .

Serializes an engine.

Parameters

engine (Union[trt.ICudaEngine, Callable() -> trt.ICudaEngine]) – An engine or a callable that returns one.

Returns

The serialized engine.

Return type

bytes

class SaveEngine(engine, path)[source]

Bases: polygraphy.backend.base.loader.BaseLoader

Functor that saves an engine to the provided path.

Saves an engine to the provided path.

Parameters
  • engine (Union[trt.ICudaEngine, Callable() -> trt.ICudaEngine]) – An engine or a callable that returns one.

  • path (str) – The path at which to save the engine.

call_impl()[source]
Returns

The engine that was saved.

Return type

trt.ICudaEngine

__call__(*args, **kwargs)

Invokes the loader by forwarding arguments to call_impl.

Note: call_impl should not be called directly - use this function instead.

save_engine(engine, path)

Immediately evaluated functional variant of SaveEngine .

Saves an engine to the provided path.

Parameters
  • engine (Union[trt.ICudaEngine, Callable() -> trt.ICudaEngine]) – An engine or a callable that returns one.

  • path (str) – The path at which to save the engine.

Returns

The engine that was saved.

Return type

trt.ICudaEngine

class OnnxLikeFromNetwork(network)[source]

Bases: polygraphy.backend.base.loader.BaseLoader

Functor that creates an ONNX-like, but not valid ONNX, model based on a TensorRT network.

[HIGHLY EXPERIMENTAL] Creates an ONNX-like, but not valid ONNX, model from a TensorRT network. This uses the ONNX format, but generates nodes that are not valid ONNX operators. Hence, this should be used only for visualization or debugging purposes.

The resulting model does not include enough information to faithfully reconstruct the TensorRT network, but does preserve the structure of the network and many of the layer parameters.

Parameters

network (Union[Tuple[trt.Builder, trt.INetworkDefinition, Optional[parser]], Callable() -> Tuple[trt.Builder, trt.INetworkDefinition, Optional[parser]]) – A tuple containing a TensorRT builder, network and optionally parser or a callable that returns one. To omit the parser, return a tuple containing just the builder and network.

call_impl()[source]
Returns

The ONNX-like, but not valid ONNX, representation of the TensorRT network.

Return type

onnx.ModelProto

__call__(*args, **kwargs)

Invokes the loader by forwarding arguments to call_impl.

Note: call_impl should not be called directly - use this function instead.

onnx_like_from_network(network)

Immediately evaluated functional variant of OnnxLikeFromNetwork .

[HIGHLY EXPERIMENTAL] Creates an ONNX-like, but not valid ONNX, model from a TensorRT network. This uses the ONNX format, but generates nodes that are not valid ONNX operators. Hence, this should be used only for visualization or debugging purposes.

The resulting model does not include enough information to faithfully reconstruct the TensorRT network, but does preserve the structure of the network and many of the layer parameters.

Parameters

network (Union[Tuple[trt.Builder, trt.INetworkDefinition, Optional[parser]], Callable() -> Tuple[trt.Builder, trt.INetworkDefinition, Optional[parser]]) – A tuple containing a TensorRT builder, network and optionally parser or a callable that returns one. To omit the parser, return a tuple containing just the builder and network.

Returns

The ONNX-like, but not valid ONNX, representation of the TensorRT network.

Return type

onnx.ModelProto