Optimization Profile

Module: polygraphy.backend.trt

class ShapeTuple(min, opt, max)[source]

Bases: object

Represents a set of shapes for a single binding in a profile.

Parameters:

  • min (Tuple[int]) – The minimum shape that the profile will support.

  • opt (Tuple[int]) – The shape for which TensorRT will optimize the engine.

  • max (Tuple[int]) – The maximum shape that the profile will support.

class Profile(dct=None)[source]

Bases: Interface

An ordered dictionary that represents a single optimization profile that can be used to build an engine.

More specifically, it is an OrderedDict[str, ShapeTuple] which maps binding names to a set of min/opt/max shapes.

add(name, min, opt, max)[source]

A convenience function to add shapes for a single binding.

Parameters:

  • name (str) – The name of the binding.

  • min (Tuple[int]) – The minimum shape that the profile will support.

  • opt (Tuple[int]) – The shape for which TensorRT will optimize the engine.

  • max (Tuple[int]) – The maximum shape that the profile will support.

Returns:

self, which allows this function to be easily chained to add multiple bindings, e.g., Profile().add(…).add(…)

Return type:

Profile

__getitem__(key)[source]

Retrieves the shapes registered for a given input name.

Returns:

A named tuple including min, opt, and max members for the shapes corresponding to the input.

Return type:

ShapeTuple

fill_defaults(network, default_shape_value=None)[source]

Fill this profile with sane default values for any bindings whose shapes have not been set explicitly.

Parameters:

  • network (trt.INetworkDefinition) – The TensorRT network this profile is meant for. This will be used to determine model inputs and their shapes.

  • default_shape_value (int) – The value to use to override dynamic dimensions.

Returns:

Self

Return type:

Profile

to_trt(builder, network)[source]

Creates a TensorRT IOptimizationProfile based on the values set in this Profile.

Parameters:

  • builder (trt.Builder) – A TensorRT builder. This will be used to construct the IOptimizationProfile.

  • network (trt.INetworkDefinition) – The TensorRT network the profile applies to.

Returns:

A TensorRT optimization profile.

Return type:

trt.IOptimizationProfile