Loaders

Module: polygraphy.tools.args

class TfTrtArgs[source]

Bases: polygraphy.tools.args.base.BaseArgs

[UNTESTED] TensorFlow-TensorRT Integration: TensorFlow-TensorRT.

Depends on:

TrtConfigArgs

parse_impl(args)[source]

Parses command-line arguments and populates the following attributes:

use_tftrt

Whether to use TF-TRT.

Type: bool

minimum_segment_size

The minimum size of segments offloaded to TRT.

Type: int

dynamic_op

Whether to enable dynamic mode, which defers engine building until runtime.

Type: bool

add_to_script_impl(script, loader_name=None, suffix=None)[source]

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

class TfLoadArgs(allow_artifacts: Optional[bool] = None, allow_custom_outputs: Optional[bool] = None, allow_tftrt: Optional[bool] = None)[source]

Bases: polygraphy.tools.args.base.BaseArgs

TensorFlow Model Loading: loading TensorFlow models.

Depends on:

ModelArgs

TfTrtArgs: if allow_tftrt == True

TrtSaveEngineBytesArgs: if allow_tftrt == True

Parameters

allow_artifacts (bool) – Whether to allow saving artifacts to the disk, like frozen models or TensorBoard visualizations. Defaults to True.
allow_custom_outputs (bool) – Whether to allow marking custom output tensors. Defaults to True.
allow_tftrt (bool) – Whether to allow applying TF-TRT. Defaults to False.

parse_impl(args)[source]

Parses command-line arguments and populates the following attributes:

ckpt

Name of the checkpoint.

Type: str

outputs

Names of output tensors.

Type: List[str]

save_frozen_graph_path

The path at which the frozen graph will be saved.

Type: str

save_tensorboard_path

The path at which the TensorBoard visualization will be saved.

Type: str

freeze_graph

Whether to attempt to freeze the graph.

Type: bool

add_to_script_impl(script, disable_custom_outputs=None)[source]

Parameters: disable_custom_outputs (bool) – Whether to disallow modifying outputs according to the outputs attribute. Defaults to False.

load_graph()[source]

Loads a TensorFlow graph according to arguments provided on the command-line.

Returns: tf.Graph

class TfConfigArgs[source]

Bases: polygraphy.tools.args.base.BaseArgs

TensorFlow Session Configuration: creating the TensorFlow SessionConfig.

parse_impl(args)[source]

Parses command-line arguments and populates the following attributes:

gpu_memory_fraction

The maximum percentage of GPU memory TensorFlow can allocate per session.

Type: float

allow_growth

Whether TensorFlow can dynamically allocate additional GPU memory.

Type: bool

xla

Whether to enable XLA.

Type: bool