Loaders¶

Module: polygraphy.tools.args

class TrtLoadPluginsArgs[source]¶

Bases: polygraphy.tools.args.base.BaseArgs

TensorRT Plugin Loading: loading TensorRT plugins.

parse_impl(args)[source]¶

Parses command-line arguments and populates the following attributes:

plugins¶

Path(s) to plugin libraries.

Type: List[str]

add_to_script_impl(script, loader_name: str)[source]¶

Parameters: loader_name (str) – The name of the loader which should be consumed by the LoadPlugins loader.

class TrtLoadNetworkArgs(allow_custom_outputs: bool = None, allow_onnx_loading: bool = None)[source]¶

Bases: polygraphy.tools.args.base.BaseArgs

TensorRT Network Loading: loading TensorRT networks.

Depends on:

ModelArgs

TrtLoadPluginsArgs

OnnxLoadArgs: if allow_onnx_loading == True

Parameters

allow_custom_outputs (bool) – Whether to allow marking custom output tensors. Defaults to True.
allow_onnx_loading (bool) – Whether to allow parsing networks from an ONNX model. Defaults to True.

parse_impl(args)[source]¶

Parses command-line arguments and populates the following attributes:

outputs¶

Names of output tensors.

Type: List[str]

exclude_outputs¶

Names of tensors which should be unmarked as outputs.

Type: List[str]

trt_network_func_name¶

The name of the function in a custom network script that creates the network.

Type: str

load_network()[source]¶

Loads a TensorRT Network model according to arguments provided on the command-line.

Returns: tensorrt.INetworkDefinition

class TrtSaveEngineArgs(output_opt: str = None, output_short_opt: str = None)[source]¶

Bases: polygraphy.tools.args.base.BaseArgs

TensorRT Engine Saving: saving TensorRT engines.

Parameters

output_opt (str) – The name of the output path option. Defaults to “output”. Use a value of False to disable the option.
output_short_opt (str) – The short option to use for the output path. Defaults to “-o”. Use a value of False to disable the short option.

parse_impl(args)[source]¶

Parses command-line arguments and populates the following attributes:

path¶

The path at which to save the TensorRT engine.

Type: str

add_to_script_impl(script, loader_name)[source]¶

Parameters: loader_name (str) – The name of the loader which should be consumed by the SaveEngine loader.
Returns: The name of the SaveEngine loader added to the script.
Return type: str

save_engine(engine, path=None)[source]¶

Saves a TensorRT engine according to arguments provided on the command-line.

Parameters

model (onnx.ModelProto) – The TensorRT engine to save.
path (str) – The path at which to save the engine. If no path is provided, it is determined from command-line arguments.

Returns

The engine that was saved.

Return type

tensorrt.ICudaEngine

class TrtLoadEngineArgs(allow_saving: bool = None)[source]¶

Bases: polygraphy.tools.args.base.BaseArgs

TensorRT Engine: loading TensorRT engines.

Depends on:

ModelArgs

TrtLoadPluginsArgs

TrtLoadNetworkArgs: if support for building engines is required

TrtConfigArgs: if support for building engines is required

TrtSaveEngineArgs: if allow_saving == True

Parameters: allow_saving (bool) – Whether to allow loaded models to be saved. Defaults to False.

parse_impl(args)[source]¶

Parses command-line arguments and populates the following attributes:

save_timing_cache¶

Path at which to save the tactic timing cache.

Type: str

add_to_script_impl(script, network_name=None)[source]¶

Parameters: network_name (str) – The name of a variable in the script pointing to a network loader.

load_engine(network=None)[source]¶

Loads a TensorRT engine according to arguments provided on the command-line.

Parameters: network (Tuple[trt.Builder, trt.INetworkDefinition, Optional[parser]]) – A tuple containing a TensorRT builder, network and optionally parser.
Returns: The engine.
Return type: tensorrt.ICudaEngine

class TrtConfigArgs(precision_constraints_default: bool = None, allow_random_data_calib_warning: bool = None, allow_custom_input_shapes: bool = None, allow_calibration: bool = None)[source]¶

Bases: polygraphy.tools.args.base.BaseArgs

TensorRT Builder Configuration: creating the TensorRT BuilderConfig.

Depends on:

ModelArgs: if allow_custom_input_shapes == True

DataLoaderArgs: if allow_calibration == True

Parameters

precision_constraints_default (str) – The default value to use for the precision constraints option. Defaults to “none”.
allow_random_data_calib_warning (bool) – Whether to issue a warning when randomly generated data is being used for calibration. Defaults to True.
allow_custom_input_shapes (bool) – Whether to allow custom input shapes when randomly generating data. Defaults to True.
allow_calibration (bool) – Whether to allow INT8 calibration. Defaults to True.

parse_impl(args)[source]¶

Parses command-line arguments and populates the following attributes:

profile_dicts¶

A list of profiles where each profile is a dictionary that maps input names to a tuple of (min, opt, max) shapes.

Type: List[OrderedDict[str, Tuple[Shape]]]

tf32¶

Whether to enable TF32.

Type: bool

fp16¶

Whether to enable FP16.

Type: bool

int8¶

Whether to enable INT8.

Type: bool

precision_constraints¶

The precision constraints to apply.

Type: str

restricted¶

Whether to enable safety scope checking in the builder.

Type: bool

calibration_cache¶

Path to the calibration cache.

Type: str

calibration_base_class¶

The name of the base class to use for the calibrator.

Type: str

sparse_weights¶

Whether to enable sparse weights.

Type: bool

load_timing_cache¶

Path from which to load a timing cache.

Type: str

load_tactics¶

Path from which to load a tactic replay file.

Type: str

save_tactics¶

Path at which to save a tactic replay file.

Type: str

tactic_sources¶

Names of the tactic sources to enable.

Type: List[str]

trt_config_script¶

Path to a custom TensorRT config script.

Type: str

trt_config_func_name¶

Name of the function in the custom config script that creates the config.

Type: str

use_dla¶

Whether to enable DLA.

Type: bool

allow_gpu_fallback¶

Whether to allow GPU fallback when DLA is enabled.

Type: bool

memory_pool_limits¶

Mapping of memory pool names to memory limits in bytes.

Type: Dict[str, int]

create_config(builder, network)[source]¶

Creates a TensorRT BuilderConfig according to arguments provided on the command-line.

Parameters

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

The TensorRT builder configuration.

Return type

trt.IBuilderConfig