转换阶段¶
一旦图被简化为易于转换的形式,我们就会设置一个转换上下文,以管理从块节点构建 TensorRT INetworkDefinition
。转换上下文记录已转换的节点集、块输入和输出以及有关图转换的其他信息。然后,此数据用于帮助转换器将图层链接在一起,并保存构建时信息,例如构建引擎所需的权重。创建上下文后,块转换器开始迭代节点列表,对于每个节点,转换器将查看其输入并组装资源数组以传递给转换器。输入可以处于以下几种状态
输入是块参数
在这种情况下,输入应已作为 IValue 存储在转换上下文
evaluated_value_map
中。转换阶段会将 IValue 添加到转换器的 args 列表中
输入是已转换节点的输出
在这种情况下,输出的 ITensor 已添加到
value_tensor_map
中。转换阶段会将 ITensor 添加到转换器的 args 列表中
输入来自生成静态值的节点
有些节点生成静态值,通常用于存储运算符的参数,我们需要在转换时评估这些节点,以便能够转换操作。转换系统将在评估器注册表中查找节点评估器,并在节点上运行它。生成的 IValue 将输入到转换上下文
evaluated_value_map
中,并添加到转换器的 args 列表中。如果要评估的节点接受输入,则转换阶段将递归地解析依赖关系,直到评估出最终的静态值
输入来自未转换的节点
Torch-TensorRT 将在此处报错
节点评估¶
有些节点包含静态数据,是操作的资源。可以在转换时评估这些节点,以便在进行节点转换时可以使用这些值。理论上,任何节点类型都可以具有转换时评估器,只要它生成静态 IValue 即可。此 IValue 将存储在转换上下文中,以便任何将评估节点作为输入的节点都可以使用它。常见的节点类型是 prim::Constant
,它发出常量,而 prim::ListConstruct
,它创建列表。
节点转换器¶
节点转换器将 JIT 节点映射到层或层子图。然后,它们将 JIT 图的输出和 TRT 图一起关联在转换上下文中。这允许转换阶段组装下一个节点的输入。在某些情况下,节点产生的输出不是 Tensor,而是对需要首先转换的输入进行计算而得到的静态结果。在这种情况下,转换器可能会将输出关联在 evaluated_value_map
而不是 value_tensor_map
中。有关更多信息,请查看:writing_converters