Conversion Phase ¶
Once the graph has be simplified to a form thats easy to convert, we then set up a conversion context
to manage the construction of a TensorRT
INetworkDefinition
from the blocks nodes. The conversion context
records the set of converted nodes, block inputs and outputs and other information about the conversion
of the graph. This data is then used to help converters link together layers and also hold build time
information like weights required to construct the engine. After the context is created, the block
converter starts iterating through the list of nodes, for each node, the converter will look at its
inputs and assemble an array of resources to pass to the converter. Inputs can be in a couple of states:
-
The input is a block parameter
-
In this case the input should have already been stored in as an IValue in the conversion context
evaluated_value_map
. The conversion stage will add the IValue to the list of args for the converter
-
-
The input is an output of a node that has already been converted
-
In this case the ITensor of the output has added to the to the
value_tensor_map
, The conversion stage will add the ITensor to the list of args for the converter
-
-
The input is from a node that produces a static value
-
There are nodes that produce static values, typically used to store parameters for operators, we need to evaluate these nodes at conversion time to be able to convert a op. The conversion system will look for a node evaluator in the evaluator registry and run it on the node. The IValue produced will be entered in the conversion context
evaluated_value_map
and added to the list of args for the converter. If the node to be evaluated takes inputs, the conversion stage will recursively resolve dependencies until the final static value has been evaluated
-
-
The input is from a node that has not been converted * Torch-TensorRT will error out here
Node Evaluation ¶
There are some nodes that contain static data and are resources for operations. These can be evaluated at
conversion time so that you can use those values when doing node conversion. In theory any node kind can have
a conversion time evaluator as long as it produces a static IValue, This IValue will be stored in the conversion
context so it can be consumed by any node that takes the evaluated node as an input. Common node types are
prim::Constant
which emits a constant and
prim::ListConstruct
which makes lists.
Node Converters ¶
Node converters map JIT nodes to layers or subgraphs of layers. They then associate outputs from the JIT graph
and the TRT graph together in the conversion context. This allows the conversion stage to assemble the inputs
for the next node. There are some cases where a node produces an output that is not a Tensor but a static result
from a calculation done on inputs which need to be converted first. In this case the converter may associate the outputs in
the
evaluated_value_map
instead of the
value_tensor_map
. For more information take a look at:
Writing Converters