Shortcuts

转换阶段

一旦图形被简化为易于转换的形式,我们就会设置一个转换上下文来管理从块节点构建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。有关更多信息,请查看:writing_converters