Loaders

Module: polygraphy.tools.args

class TfTrtArgs[source]

Bases: 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: bool | None = None, allow_custom_outputs: bool | None = None, allow_tftrt: bool | None = None)[source]

Bases: 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: 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