Loaders

Module: polygraphy.tools.args

class TfTrtArgs

Bases: polygraphy.tools.args.base.BaseArgs

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

Depends on:

• TrtConfigArgs

• TrtLegacyRunnerArgs

parse_impl(args)

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)

Parameters

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

class TfLoadArgs(allow_artifacts: bool = None, allow_custom_outputs: bool = None, allow_tftrt: bool = None)

Bases: polygraphy.tools.args.base.BaseArgs

TensorFlow Model Loading: loading TensorFlow models.

Depends on:

• ModelArgs

• TfTrtArgs: if allow_tftrt == True

• TrtSaveEngineArgs: 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)

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)

Parameters

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

load_graph()

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

Returns

tf.Graph

class TfConfigArgs

Bases: polygraphy.tools.args.base.BaseArgs

TensorFlow Session Configuration: creating the TensorFlow SessionConfig.

parse_impl(args)

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