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