Algorithm Selector

Module: polygraphy.backend.trt

class TensorInfo(dtype, strides, vectorized_dim, components_per_element)[source]

Bases: object

Tracks information about a tensor, such as format and data type.

Parameters:

dtype (trt.DataType) – The data type.
strides (Sequence[int]) – The strides.
vectorized_dim (int) – The index of the vectorized dimensions.
components_per_element (int) – The number of components per element.

static from_trt(io_info)[source]

Creates a Polygraphy TensorInfo instance from a TensorRT IAlgorithmIOInfo.

Parameters:: io_info (trt.IAlgorithmIOInfo) – The algorithm I/O information.
Returns:: TensorInfo

class Algorithm(implementation, tactic, inputs, outputs)[source]

Bases: object

Represents a TensorRT algorithm variant, which can be uniquely represented by an implementation ID, tactic ID, and I/O tensor information.

Parameters:

implementation (int) – The implementation for this Algorithm.
tactic (int) – The tactic for this Algorithm.
inputs (Sequence[TensorInfo]) – A sequence of TensorInfos for each input.
outputs (Sequence[TensorInfo]) – A sequence of TensorInfos for each output.

static from_trt(context, algorithm)[source]

Creates a Polygraphy Algorithm instance from a TensorRT IAlgorithmContext and IAlgorithm.

Parameters:

context (trt.IAlgorithmContext) – The algorithm context corresponding to the layer.
algorithm (trt.IAlgorithm) – The algorithm variant provided by TensorRT.

Returns:

Algorithm

class TacticReplayData(dct=None)[source]

Bases: Interface

Maps layer names to corresponding tactics. More specifically, it is an OrderedDict[str, Algorithm].

add(name, algorithm)[source]

Add an entry into the tactic replay data.

Parameters:

name (str) – The name of the layer
algorithm (Algorithm) – The algorithm to use for the layer.

Returns:

self, to allow for method chaining.

Return type:

TacticReplayData

static from_json(src)

Decode a JSON object and create an instance of this class.

Parameters:: src (str) – The JSON representation of the object
Returns:: The decoded instance
Return type:: TacticReplayData
Raises:: PolygraphyException – If the JSON cannot be decoded to an instance of TacticReplayData

static load(src)

Loads an instance of this class from a JSON file.

Parameters:: src (Union[str, file-like]) – The path or file-like object to read from.
Returns:: The decoded instance
Return type:: TacticReplayData
Raises:: PolygraphyException – If the JSON cannot be decoded to an instance of TacticReplayData

save(dest)

Encode this instance as a JSON object and save it to the specified path or file-like object.

Parameters:: dest (Union[str, file-like]) – The path or file-like object to write to.

to_json()

Encode this instance as a JSON object.

Returns:: A JSON representation of this instance.
Return type:: str

TacticRecorder(record)[source]

A TensorRT algorithm selector that can record tactics selected by TensorRT.

The generated tactic replay file is specific to network and builder configuration. Changing either of these may render the tactic replay file unusable.

Parameters:: record (Union[path, file-like, TacticReplayData]) – A path or file-like object or an empty TacticReplayData instance. Tactics will be recorded and stored here.

TacticReplayer(replay)[source]

A TensorRT algorithm selector that can replay tactics according to a tactic replay file.

Parameters:: replay (Union[path, file-like, TacticReplayData]) – A path or file-like object containing a JSON-ified TacticReplayData instance, or a TacticReplayData instance.