Algorithm Selector

Module: polygraphy.backend.trt

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

Bases: object

Represents a TensorRT algorithm variant, which can be uniquely represented by an implementation ID and tactic ID.

Parameters

implementation (int) – The implementation for this Algorithm.
tactic (int) – The tactic for this Algorithm.
inputs (List[Tuple[trt.TensorFormat, trt.DataType, Sequence[int]]]) – A list of tuples containg a TensorRT tensor format, data type, and strides for each input.
outputs (List[Tuple[trt.TensorFormat, trt.DataType, Sequence[int]]]) – A list of tuples containg a TensorRT tensor format, data type, and strides 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.

class TacticReplayData(dct=None)[source]

Bases: polygraphy.common.interface.TypedDict.<locals>.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.