转换阶段

一旦图被简化为易于转换的形式，我们就会设置一个转换上下文来管理从块节点构建 TensorRT INetworkDefinition。转换上下文记录已转换节点、块输入和输出以及有关图转换的其他信息。然后，这些数据用于帮助转换器将层链接在一起，并保存构建时信息，例如构建引擎所需的权重。创建上下文后，块转换器开始遍历节点列表，对于每个节点，转换器会查看其输入并组装一个资源数组以传递给转换器。输入可以处于几种状态：

• 输入是一个块参数

  • 在这种情况下，输入应已作为 IValue 存储在转换上下文的 evaluated_value_map 中。转换阶段会将 IValue 添加到转换器的参数列表中。

• 输入是已转换节点的输出

  • 在这种情况下，输出的 ITensor 已添加到 value_tensor_map 中。转换阶段会将 ITensor 添加到转换器的参数列表中。

• 输入来自生成静态值的节点

  • 有些节点会生成静态值，通常用于存储操作符的参数，我们需要在转换时评估这些节点才能转换操作符。转换系统会在评估器注册表中查找节点评估器并对其运行。生成的 IValue 将输入到转换上下文的 evaluated_value_map 中，并添加到转换器的参数列表中。如果要评估的节点接受输入，转换阶段将递归地解析依赖关系，直到最终静态值被评估。

• 输入来自尚未转换的节点

  • Torch-TensorRT 在此处会报错

节点评估

有些节点包含静态数据，并且是操作的资源。这些节点可以在转换时进行评估，以便在进行节点转换时可以使用这些值。理论上，任何节点类型都可以有一个转换时评估器，只要它产生一个静态 IValue，这个 IValue 将存储在转换上下文中，以便任何将评估节点作为输入的节点都可以使用它。常见的节点类型是 prim::Constant（发出常量）和 prim::ListConstruct（创建列表）。

节点转换器

节点转换器将 JIT 节点映射到层或层子图。然后，它们在转换上下文中关联 JIT 图和 TRT 图的输出。这允许转换阶段为下一个节点组装输入。在某些情况下，节点产生的输出不是张量，而是对输入进行计算的静态结果，这些输入需要先进行转换。在这种情况下，转换器可能会在 evaluated_value_map 中而不是 value_tensor_map 中关联输出。欲了解更多信息，请参阅：编写转换器