编写 TorchScript 转换器¶

背景¶

在 JIT IR 中，操作表示为图中的节点。节点有输入和输出，由 torch::jit::Values 表示，它们是流入和流出节点的数据的类型化抽象表示。TensorRT 通过使用 nvinfer1::ILayers 和 nvinfer1::ITensors 来表示其图，它们分别对应于节点和值。转换器的目标是创建新的 ILayers 和子图，以执行节点指定的操作，并将生成的 ITensors 和 Values 关联起来。

转换器¶

转换器应该是函数，它们将使用输入列表（可以是 nvinfer1::ITensors 或 torch::jit::IValues）来构建与 LibTorch 操作等效的层。

转换器可以使用 RegisterNodeConversionPatterns 辅助类进行注册，您实例化一个 RegisterNodeConversionPatterns 对象，并调用其上的 pattern 函数（如下所示），该函数接受一个字符串，描述将触发转换器运行的操作的函数 schema，以及一个执行实际转换的 lambda 或函数。

请注意，pattern 函数可以链式调用。

auto acthardtanh TORCHTRT_UNUSED = RegisterNodeConversionPatterns()
    .pattern({
        "aten::hardtanh(Tensor self, Scalar min_val=-1, Scalar max_val=1) -> (Tensor)",
        [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
            auto in = args[0].ITensor();
            auto min = args[1].unwrapToDouble();
            auto max = args[2].unwrapToDouble();

            auto new_layer = ctx->net->addActivation(*in, nvinfer1::ActivationType::kCLIP);
            TORCHTRT_CHECK(new_layer, "Unable to create layer for aten::hardtanh");

            new_layer->setAlpha(min);
            new_layer->setBeta(max);

            new_layer->setName(util::node_info(n).c_str());
            auto out_tensor = ctx->AssociateValueAndTensor(n->outputs()[0], new_layer->getOutput(0));

            LOG_DEBUG("Output shape: " << out_tensor->getDimensions());
            return true;
        }
    });

转换器约定¶

对转换器的保证¶

在 args 中，对于每个节点输入值，都会有一个条目，可以是 ITensor 或 IValue。
输入将按照函数 schema 的顺序提供。

转换器的职责¶

必须保证 args 是可以直接解包 Arg 联合类型的特定类型，无需检查；通常，输入张量参数可以预期为 ITensors。
必须保证所有权重或静态值在转换结束前始终有效。
1. 一个有用的工具是下文描述的 Weights 辅助类。
转换器应为节点的每个输出生成一个 IValue 或 ITensor。如果存在没有关联的 ITensors 或 IValues 的 Values，编译器将进行检查并发出警告。
输出必须进行标注。
1. 在转换上下文的 value_tensor_map 中，JIT 节点的输出值与新的 TRT 层的输出张量之间必须建立关联。
为您的层命名。
1. 当我们跟踪哪些层和节点相互对应时，调试会容易得多。我们目前使用的系统是使用节点的“节点信息”作为层的名称。
为您的张量命名。
1. 使用输出值的调试名称作为新 ITensor 的名称（同样为了调试）。

转换上下文¶

转换上下文维护转换状态，它管理网络定义、两个映射（一个存储 Values 和 IValues 之间的关联，即 evaluated_value_map；另一个存储 Values 和 ITensors 之间的关联）以及需要在转换结束前存活的任何类型的内存。您在转换器中主要会交互的 API 是直接访问网络定义以添加层（ctx->net）和数据关联函数（ctx->AssociateValueAndTensor() 和 ctx->AssociateValueAndIValue()），您将使用这些 API 将层添加到 TRT 层，并记录节点输出与静态值或 TensorRT 层输出的配对。

参数¶

提供给转换器的参数是 nvinfer1::ITensors 和 torch::jit::IValues 的可检查联合类型（即 TensorRT 图中的抽象数据流和静态值）。对于节点的每个输入值，您都可以保证会获得相应的参数。它们按照函数 schema 的顺序提供。可以预期输入（即在 PyTorch 中传递给模块 forward 函数的参数）将是 ITensors，但如果您不确定，Arg 类也有机制在解包前安全地检查参数。Arg 类还具有深度解包方法，如果您确定安全，可以直接获取 IValue 中的底层数据。如果 IValue 可能为 None，您还可以传入一个备用值。IValues 已被扩展，仅在 TensorLists 的情况下能够包含围绕 ITensors 的包装器。您可以使用类似于以下模式从 IValue 获取 ITensor：ivalue.toCustomClass<TensorContainer>()->tensor()。您可以使用 ivalue.isTensor() 或 ivalue.isCustomClass() 判断 IValue 是否包含 Tensor 或 ITensor。

权重¶

权重在构建时使用，因此必须保证所有权重在转换阶段结束前始终存活。TensorRT 也有自己的权重结构来保存权重。转换器可以使用一个围绕此类的包装器，该包装器抽象了许多细节。

权重包装类可以接受 at::Tensors 或单个值（目前）。在构建这些权重时，您还需要传入转换上下文，因为在内部，权重类将分配由转换上下文管理的内存来存储张量数据的副本。当转换上下文的析构函数被销毁时，这些数据就会被释放，因此转换器无需真正考虑它。

输入数据的形状会生成元数据，这在与 TensorRT 交互时非常有用，例如输入映射数量、输出映射数量和内核形状。

其他建议¶

在处理权重和其他静态值时，您可以受益于完整的 aten 库。这意味着您可以在转换期间进行大量工作以实现高效转换。一个很好的例子是 batch_norm 转换器，它在创建 TensorRT 层之前与 PyTorch 一起执行操作融合。