编写转换器¶

背景¶

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

转换器¶

转换器应该是使用输入列表（nvinfer1::ITensors 或 torch::jit::IValues）来构建等效于 LibTorch 运算的层的函数。

转换器可以使用 RegisterNodeConversionPatterns 辅助类进行注册，其中你实例化一个 RegisterNodeConversionPatterns 对象并对其调用 pattern 函数（如下所示），该函数接受一个字符串，描述将导致转换器运行的操作的函数模式，以及一个 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。
输入将根据函数模式的顺序提供。

转换器的责任¶

必须保证 Arg 是一种类型，可以在不进行检查的情况下解包 Arg 联合，通常可以预期输入张量参数为 ITensor。
必须保证任何权重或静态值在转换结束之前有效。
1. 一个有用的工具是下面描述的 Weights 辅助类。
转换器应为节点的每个输出生成 IValue 或 ITensor。编译器将检查这一点，如果存在没有关联 ITensor 或 IValue 的 Value，则会产生警告。
输出必须进行注释。
1. 在转换上下文中，必须在 JIT 节点的输出值和新的 TRT 层的输出张量之间建立关联，即 value_tensor_map。
命名你的层
1. 当我们可以跟踪哪些层和节点对应于彼此时，调试会容易得多。我们目前使用的系统是使用节点的“节点信息”作为层的名称。
命名你的张量
1. 使用输出值的调试名称作为新 ITensor 的名称（同样用于调试）。

转换上下文¶

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

参数¶

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

权重¶

权重在构建期间使用，因此需要保证任何权重在转换阶段结束之前都存在。TensorRT 还使用它自己的权重结构来保存权重。有一个转换器可用的权重包装器类，它抽象化了这些内容的大部分。

权重包装器类可以接受 at::Tensors 或单个值（目前如此）。你还需要在构造这些权重时传递转换上下文，因为权重类在内部会使用转换上下文管理的内存来存储张量数据的副本。当转换上下文析构函数被销毁时，这些数据将被释放，因此转换器实际上不需要考虑它。

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

其他建议¶

在处理权重和其他静态值时，你可以使用完整的 aten 库。这意味着你可以在转换期间做相当多的工作来生成高效的转换。一个很好的例子是 batch_norm 转换器，其中转换器在创建 TensorRT 层之前使用 PyTorch 对操作进行融合。