快捷方式

转换阶段

图简化为易于转换的形式后,我们将设置一个转换上下文,用于管理从块节点构建 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 图的输出关联起来。这使得转换阶段能够为下一个节点组装输入。在某些情况下,节点产生的输出不是 Tensor,而是对需要首先转换的输入进行计算产生的静态结果。在这种情况下,转换器可能会将输出关联到 evaluated_value_map 中,而不是 value_tensor_map 中。更多信息请参阅:writing_converters

文档

访问 PyTorch 的全面开发者文档

查看文档

教程

获取面向初学者和高级开发者的深入教程

查看教程

资源

查找开发资源并获得问题解答

查看资源