编写TorchScript转换器¶

背景¶

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

转换器¶

转换器应该是使用输入列表（无论是nvinfer1::ITensors还是torch::jit::IValues）来构建与LibTorch操作等效的层的函数。

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

注意模式函数可以链式调用

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;
        }
    });

转换器合约¶

转换器保证什么¶

在参数中，每个节点输入值都会有一个条目，无论是ITensor还是IValue
输入将根据函数模式按顺序提供

转换器的职责¶

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

转换上下文¶

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

参数¶

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

权重¶

权重在构建时使用，因此需要确保任何权重在转换阶段结束之前一直存在。 TensorRT 还使用其自己的权重结构来保存权重。有一个围绕此类的包装器可用，用于转换，它抽象了很多这方面的内容。

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

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

其他建议¶

在处理权重和其他静态值时，您可以充分利用完整的aten库。这意味着您可以在转换期间进行大量工作，以生成高效的转换。一个很好的例子是batch_norm转换器，在创建TensorRT层之前，转换器会与PyTorch进行操作的融合。