使用自定义层扩展 TensorRT#

NVIDIA TensorRT 支持许多层，并且其功能在不断扩展；但是，在某些情况下，支持的层可能无法满足模型的特定需求。在这种情况下，可以通过实现自定义层（通常称为插件）来扩展 TensorRT。

TensorRT 包含可以加载到您的应用程序中的标准插件。有关开源插件的列表，请参阅 GitHub: TensorRT 插件。

要在您的应用程序中使用标准 TensorRT 插件，必须加载 libnvinfer_plugin.so（Windows 上为 nvinfer_plugin.dll）库，并且必须通过在您的应用程序代码中调用 initLibNvInferPlugins 来注册所有插件。有关这些插件的更多信息，请参阅 NvInferPlugin.h 文件。

如果这些插件不能满足您的需求，您可以编写并添加自己的插件。

使用 C++ API 添加自定义层#

确保 TensorRT 正确识别您的插件需要四个步骤

  1. 从 TensorRT 的插件基类之一实现插件类。目前，唯一推荐的是 IPluginV3。

  2. 通过从 TensorRT 的基于插件创建器的类之一派生，实现与您的类绑定的插件创建器类。目前，唯一推荐的是 IPluginCreatorV3One。

  3. 在 TensorRT 的插件注册表中注册插件创建器类的实例。

  4. 通过直接使用 TensorRT 的网络 API 或使用 TensorRT ONNX 解析器 API 加载 ONNX 模型，将插件类的实例添加到 TensorRT 网络。

以下部分详细探讨了这些步骤中的每一个步骤。

实现插件类#

您可以通过从 TensorRT 的插件基类之一派生来实现自定义层。从 TensorRT 10.0 开始，唯一推荐的插件接口是 IPluginV3，因为其他接口已弃用。因此，本节主要描述使用 IPluginV3 实现插件。有关如何将实现 V2 插件接口的插件迁移到 IPluginV3 的信息，请参阅 将 V2 插件迁移到 IPluginV3 部分。

IPluginV3 是定义三个功能capability interfaces的包装器：核心、构建和运行时。

  • 核心功能：指插件属性和行为，这些属性和行为在插件生命周期的构建和运行时阶段都很常见。

  • 构建功能：指插件必须为 TensorRT 构建器展示的插件属性和行为。

  • 运行时功能：指插件必须展示的插件属性和行为，以便在 TensorRT 构建阶段的自动调优或 TensorRT 运行时阶段的推理期间可执行。

IPluginV3OneCore (C++, Python), IPluginV3OneBuild (C++, Python), 和 IPluginV3OneRuntime (C++, Python) 是 IPluginV3 插件必须实现的基类，分别用于显示核心、构建和运行时功能。如果需要 I/O 别名，则可以使用 IPluginV3OneBuildV2 (C++, Python) 作为构建功能，其中包含 IPluginV3OneBuild 中功能的超集。

实现插件创建器类#

要在网络中使用插件，您必须首先在 TensorRT 的 PluginRegistry (C++, Python) 中注册它。您不是直接注册插件，而是注册插件的工厂类的实例，该实例派生自 IPluginCreatorInterface (C++, Python) 的子类。插件创建器类还提供有关插件的其他信息：其名称、版本和插件字段参数。

IPluginCreatorV3One 是 IPluginV3 的工厂类。IPluginCreatorV3One::createPlugin()（具有以下签名）返回类型为 IPluginV3 的插件对象。

1IPluginV3* createPlugin(AsciiChar const *name, PluginFieldCollection const *fc, TensorRTPhase phase)

1create_plugin(self: trt.IPluginCreatorV3, name: str, field_collection: trt.PluginFieldCollection, phase: trt.TensorRTPhase) -> trt.IPluginV3

可以调用 IPluginCreatorV3One::createPlugin() 以在 TensorRT 的构建阶段或 TensorRT 的运行时阶段创建插件实例，这通过 phase 参数（类型为 TensorRTPhase (C++, Python)）进行通信。

  • 返回的 IPluginV3 对象在两个阶段都必须具有有效的核心功能。

  • 在构建阶段，返回的 IPluginV3 对象必须同时具有构建和运行时功能。

  • 在运行时阶段，返回的 IPluginV3 对象必须具有运行时功能。构建功能不是必需的，将被忽略。

在插件注册表中注册插件创建器#

有两种方法可以在注册表中注册插件

  1. TensorRT 提供了一个宏 REGISTER_TENSORRT_PLUGIN，该宏静态地在注册表中注册插件创建器。REGISTER_TENSORRT_PLUGIN 始终在默认命名空间（“”）下注册创建器。

  2. 通过创建类似于 initLibNvInferPlugins 的入口点并在插件注册表上调用 registerCreator 来动态注册。与静态注册相比，这更可取，因为它允许在唯一的命名空间下注册插件。这确保了在不同插件库的构建期间不会发生名称冲突。

在序列化期间，TensorRT 引擎在内部存储所有插件的插件名称、插件版本和命名空间（如果存在），以及 IPluginV3OneRuntime::getFieldsToSerialize() 返回的 PluginFieldCollection 中的任何插件字段。在反序列化期间，TensorRT 从插件注册表中查找具有相同插件名称、版本和命名空间的插件创建器，并在其上调用 IPluginCreatorV3One:::createPlugin()——序列化的 PluginFieldCollection 作为 fc 参数传递回来。

将插件实例添加到 TensorRT 网络#

您可以使用 addPluginV3() 将插件添加到 TensorRT 网络，这将创建一个具有给定插件的网络层。

例如，您可以按如下方式将插件层添加到您的网络

// Look up the plugin in the registry
// Cast to appropriate child class of IPluginCreatorInterface
auto creator = static_cast<IPluginCreatorV3One*>(getPluginRegistry()->getCreator(pluginName, pluginVersion, pluginNamespace));
PluginFieldCollection const* pluginFC = creator->getFieldNames();
// Populate the field parameters for the plugin layer
// PluginFieldCollection *pluginData = parseAndFillFields(pluginFC, layerFields);
// Create the plugin object using the layerName and the plugin metadata for use by the TensorRT builder
IPluginV3 *pluginObj = creator->createPlugin(layerName, pluginData, TensorRTPhase::kBUILD);
// Add the plugin to the TensorRT network
auto layer = network.addPluginV3(inputs.data(), int(inputs.size()),  shapeInputs.data(), int(shapeInputs.size()), pluginObj);
… (build rest of the network and serialize engine)
// Delete the plugin object
delete pluginObj;
… (free allocated pluginData)

先前描述的 createPlugin 方法在堆上创建一个新的插件对象并返回一个指针。如前所示，请确保删除 pluginObj 以避免内存泄漏。

当引擎被删除时，引擎会销毁在构建期间创建的插件对象的任何克隆。您有责任确保在将插件对象添加到网络后将其释放。

注意

  • 不要序列化所有插件参数，仅序列化运行时正确运行所需的参数。可以省略构建时参数。

  • 如果您是汽车安全用户，则必须调用 getSafePluginRegistry() 而不是 getPluginRegistry()。您还必须使用宏 REGISTER_SAFE_TENSORRT_PLUGIN 而不是 REGISTER_TENSORRT_PLUGIN。

示例：使用 C++ 添加具有动态形状的自定义层#

假设需要一个自定义层来进行类似填充的操作，其中输入批次中的每个图像都必须重塑为 32 x 32。输入张量 X 的形状为 (B, C, H, W)，输出 Y 的形状为 (B, C, 32, 32)。为了实现这一点，可以使用 IPluginV3 接口编写 TensorRT 插件；我们将其称为 PadPlugin。

由于 IPluginV3 插件必须具有多个功能，每个功能由单独的接口定义，因此您可以使用组合或多重继承的原则来实现插件。但是，对于大多数用例来说，多重继承方法更容易，尤其是在将构建和运行时功能耦合在单个类中是可容忍的情况下。

使用多重继承，PadPlugin 可以按如下方式实现

class PadPlugin : public IPluginV3, public IPluginV3OneCore, public IPluginV3OneBuild, public IPluginV3OneRuntime
{
    ...override inherited virtual methods.
};

IPluginV3::getCapabilityInterface 的重写必须返回指向各个功能接口的指针。对于每个 PluginCapabilityType，必须通过相应的功能接口进行强制转换，以消除编译器的歧义。

IPluginCapability* PadPlugin::getCapabilityInterface(PluginCapabilityType type) noexcept override
{
    TRY
    {
        if (type == PluginCapabilityType::kBUILD)
        {
            return static_cast<IPluginV3OneBuild*>(this);
        }
        if (type == PluginCapabilityType::kRUNTIME)
        {
            return static_cast<IPluginV3OneRuntime*>(this);
        }
        ASSERT(type == PluginCapabilityType::kCORE);
        return static_cast<IPluginV3OneCore*>(this);
    }
    CATCH
    {
        // log error
    }
    return nullptr;
}

在此特定示例中，重要的方法是

  • INetworkDefinition::addPluginV3

  • IPluginV3OneBuild::getNbOutputs

  • IPluginV3OneBuild::getOutputDataTypes

  • IPluginV3OneBuild::getOutputShapes

  • IPluginV3OneBuild::supportsFormatCombination

  • IPluginV3OneBuild::configurePlugin

  • IPluginV3OneRuntime::onShapeChange

  • IPluginV3OneRuntime::enqueue

INetworkDefinition::addPluginV3 (C++, Python) 可以将插件添加到网络。

std::vector<ITensor*> inputs{X};

auto pluginLayer = network->addPluginV3(inputs.data(), inputs.size(), nullptr, 0, *plugin);

您可以通过重写 IPluginV3OneBuild::getNbOutputs 来传达存在单个插件输出。

int32_t PadPlugin::getNbOutputs() const noexcept override
{
    return 1;
}

输出将与输入具有相同的数据类型，这可以在 IPluginV3OneBuild::getOutputDataTypes 的重写中传达。

int32_t PadPlugin::getOutputDataTypes(
        DataType* outputTypes, int32_t nbOutputs, DataType const* inputTypes, int32_t nbInputs) const noexcept override
{
    outputTypes[0] = inputTypes[0];
    return 0;
}

getOutputShapes 的重写返回输出维度的符号表达式，以输入维度表示，但在数据相关的输出形状的情况下除外，这将在稍后的 示例：使用 C++ 添加具有数据相关和形状输入相关形状的自定义层 中介绍。在当前示例中，输出的前两个维度将分别等于输入的前两个维度，最后两个维度将是常数，每个维度都等于 32。IExprBuilder 传递到 getOutputShapes 中可用于定义常量符号表达式。

int32_t PadPlugin::getOutputShapes(DimsExprs const* inputs, int32_t nbInputs, DimsExprs const* shapeInputs, int32_t nbShapeInputs, DimsExprs* outputs, int32_t nbOutputs, IExprBuilder& exprBuilder) noexcept
{
    outputs[0].nbDims = 4;
    // first two output dims are equal to the first two input dims
    outputs[0].d[0] = inputs[0].d[0];
    outputs[0].d[1] = inputs[0].d[1];
    // The last two output dims are equal to 32
    outputs[0].d[2] = exprBuilder.constant(32);
    outputs[0].d[3] = exprBuilder.constant(32);
    return 0;
}

TensorRT 使用 supportsFormatCombination 来询问插件是否接受给定位置 pos 的连接的给定类型和格式组合，以及索引较小的连接的给定格式/类型。接口统一索引输入/输出作为连接，从第一个输入的 0 开始，然后按顺序排列其余输入，然后编号输出。在示例中，输入是连接 0，输出是连接 1。

为了简单起见，该示例仅支持线性格式和 FP32 类型。

bool PadPlugin::supportsFormatCombination(
        int32_t pos, DynamicPluginTensorDesc const* inOut, int32_t nbInputs, int32_t nbOutputs) noexcept override
{
    assert(0 <= pos && pos < 2);
    return inOut[pos].desc.format == PluginFormat::kLINEAR && inOut[pos].desc.type == DataType::kFLOAT;
}

TensorRT 调用两个方法，允许插件在 enqueue() 之前做出任何配置选择，无论是在自动调优期间（在引擎构建阶段）还是在引擎执行时（在运行时阶段）。

  1. IPluginV3OneBuild::configurePlugin：当插件正在准备进行分析（自动调优）但不是针对任何特定输入大小时调用。DynamicPluginTensorDesc 的 min、max 和 opt 值对应于张量形状的边界及其用于自动调优的形状。desc.dims 字段对应于在网络创建时指定的插件维度，包括动态维度的任何通配符 (-1)。

  2. IPluginV3OneRuntime::onShapeChange：在 enqueue() 之前在构建阶段和运行时阶段调用，以传达后续 enqueue() 的输入和输出形状。输出 PluginTensorDesc 将包含通过 getOutputShapes() 指定的任何数据相关维度的通配符 (-1)。

此插件不需要 configurePlugin 和 onShapeChange 做任何事情，因此它们都是空操作

int32_t PadPlugin::configurePlugin(DynamicPluginTensorDesc const* in, int32_t nbInputs, DynamicPluginTensorDesc const* out, int32_t nbOutputs) noexcept override
{
    return 0;
}

int32_t PadPlugin::onShapeChange(PluginTensorDesc const* in, int32_t nbInputs, PluginTensorDesc const* out, int32_t nbOutputs) noexcept override
{
    return 0;
}

最后，重写 PadPlugin::enqueue 必须完成工作。由于形状是动态的，因此 enqueue 会收到一个 PluginTensorDesc，它描述了每个输入和输出的维度、类型和格式。

int32_t enqueue(PluginTensorDesc const* inputDesc, PluginTensorDesc const* outputDesc, void const* const* inputs,
        void* const* outputs, void* workspace, cudaStream_t stream) noexcept override
{
    // populate outputs and return status code
}

示例：使用 C++ 添加具有数据相关和形状输入相关形状的自定义层#

本节展示了一个具有数据相关和形状输入相关形状的插件示例。请注意，数据相关的输出形状和向插件添加形状输入是 V2 插件中不存在的新功能。

  • 数据相关形状 (DDS)：插件输出的形状可能取决于输入张量的值。

  • 形状输入：插件可以接受形状和设备张量输入。这些输入仅对插件可见，作为 IPluginV3OneBuild::getOutputShapes() 的参数。因此，它们的唯一目的是帮助插件执行输出形状计算。

例如，BarPlugin 是一个具有一个设备输入 X、一个形状输入 S 和一个输出 Y 的插件，其中

  • Y 的第一个维度取决于 S 的值。

  • Y 的第二个维度是静态的。

  • Y 的第三个维度取决于 X 的形状。

  • Y 的第四个维度是数据相关的。

与先前示例中的 PadPlugin 类似，BarPlugin 使用多重继承。

要将插件添加到网络，可以使用 INetworkDefinition::addPluginV3 (C++, Python) 以类似的方式使用。在设备张量输入之后，addPluginV3 采用两个额外的参数来指定形状张量输入。

std::vector<ITensor*> inputs{X};
std::vector<ITensor*> shapeInputs{S};

auto pluginLayer = network->addPluginV3(inputs.data(), inputs.size(), shapeInputs.data(), shapeInputs.size(), *plugin);

注意

TensorRT ONNX 解析器提供了一个内置功能，可以将形状输入传递给基于 IPluginV3 的插件支持的自定义操作。要解释为形状输入的输入的索引必须由名为 tensorrt_plugin_shape_input_indices 的节点属性指示为整数列表。例如，如果自定义操作有四个输入，并且第二个和第四个输入应作为形状输入传递给插件，请添加一个名为 tensorrt_plugin_shape_input_indices 的节点属性，其类型为 onnx.AttributeProto.ints，包含值 [1, 3]。

在 getOutputShapes 的重写中，插件必须声明每个输出张量的每个数据相关维度的位置和边界。可以使用称为大小张量的特殊输出表示边界。

大小张量是 INT32 或 INT64 数据类型的标量，通过自动调优的值和上限来表示；这些值可以是常量，也可以使用 IExprBuilder 根据设备输入形状或形状输入值计算。

在这种情况下，存在一个奇异的数据相关维度，我们可以使用一个大小张量来表示。请注意，表示数据相关维度所需的任何大小张量都算作插件的输出；因此，插件总共将有两个输出。

int32_t getNbOutputs() const noexcept override
{
    return 2;
}

假设输出 Y 与设备输入 X 的类型相同，并且数据相关维度大小适合 INT32（大小张量的类型为 r）。然后 BarPlugin 像这样表示输出数据类型

int32_t getOutputDataTypes(
        DataType* outputTypes, int32_t nbOutputs, DataType const* inputTypes, int32_t nbInputs) const noexcept override
{
    outputTypes[0] = inputTypes[0];
    outputTypes[1] = DataType::kINT32;
    return 0;
}

方法 getOutputShapes 可以使用传递给它的 IExprBuilder 构建符号输出形状表达式。在下面的内容中，请注意，大小张量必须显式声明为 0D。

int32_t BarPlugin::getOutputShapes(DimsExprs const* inputs, int32_t nbInputs, DimsExprs const* shapeInputs, int32_t nbShapeInputs, DimsExprs* outputs, int32_t nbOutputs, IExprBuilder& exprBuilder) noexcept
{
    outputs[0].nbDims = 4;
    // The first output dimension depends on the value of S.
    // The value of S is encoded as fictitious dimensions.
    outputs[0].d[0] = shapeInputs[0].d[0];
    // The third output dimension depends on the shape of X
    outputs[0].d[2] = inputs[0].d[0];
    // The second output dimension is static
    outputs[0].d[1] = exprBuilder.constant(3);

    auto upperBound = exprBuilder.operation(DimensionOperation::kPROD, *inputs[0].d[2], *inputs[0].d[3]);
    auto optValue = exprBuilder.operation(DimensionOperation::kFLOOR_DIV, *upperBound, *exprBuilder.constant(2));

    // output at index 1 is a size tensor
    outputs[1].nbDims = 0; // size tensors must be declared as 0-D
    auto sizeTensor = exprBuilder.declareSizeTensor(1, *optValue, *upperBound);

    // The fourth output dimension is data-dependent
    outputs[0].d[3] = sizeTensor;

    return 0;
}

supportsFormatCombination 的重写施加以下条件

  • 设备输入 X 必须具有 DataType::kFLOAT 或 DataType::kHALF。

  • 输出 Y 必须与 X 的类型相同。

  • 大小张量输出的类型为 DataType::kINT32。

注意

通过 addPluginV3 (C++, Python) 传递给插件的形状输入仅作为 getOutputShapes() 的参数出现，并且不计算在内，也不包含在任何其他插件接口方法中的插件输入中。

bool BarPlugin::supportsFormatCombination(
        int32_t pos, DynamicPluginTensorDesc const* inOut, int32_t nbInputs, int32_t nbOutputs) noexcept override
    {
        assert(0 <= pos && pos < 3);
        auto const* in = inOut;
        auto const* out = inOut + nbInputs;

        bool typeOk{false};

        switch (pos)
        {
        case 0: typeOk = in[0].desc.type == DataType::kFLOAT || in[0].desc.type == DataType::kHALF; break;
        case 1: typeOk = out[0].desc.type == in[0].desc.type; break;
        case 2: typeOk = out[1].desc.type == DataType::kINT32; break;
        }

        return inOut[pos].desc.format == PluginFormat::kLINEAR && typeOk;
    }

这里的局部变量 in 和 out 允许按输入或输出编号而不是连接编号来检查 inOut。

重要提示

重写检查索引小于 pos 的连接的格式/类型，但绝不能检查索引大于 pos 的连接的格式/类型。示例使用 case 1 来检查 connection 1 与 connection 0，而不是使用 case 0 来检查 connection 0 与 connection 1。

configurePlugin 和 onShapeChange 在这里也将是空操作；需要注意的一件事是，在 onShapeChange 中，输出的 PluginTensorDesc 将包含数据相关维度的通配符 (-1)。

实现具有数据相关输出形状的 enqueue 与静态或动态形状情况大相径庭。与任何其他输出一样，对于具有数据相关维度的输出，传递给 enqueue 的输出缓冲区保证足够大，可以容纳相应的输出张量（基于通过 getOutputShapes 指定的上限）。

示例：使用 C++ 添加具有 INT8 I/O 支持的自定义层#

PoolPlugin 是一个插件，演示了如何使用 IPluginV3 为自定义池化层 addINT8 I/O。PoolPlugin 多重继承自 IPluginV3、IPluginV3OneCore、IPluginV3OneBuild 和 IPluginV3OneRuntime，类似于上面的 PadPlugin 和 BarPlugin 示例。

影响 INT8 I/O 的主要方法是

  • supportsFormatCombination

  • configurePlugin

supportsFormatCombination 的重写必须指示允许哪些 INT8 I/O 组合。此接口类似于 示例：使用 C++ 添加具有动态形状的自定义层。在此示例中，支持的 I/O 张量格式是具有 FP32、FP16、BF16、FP8 或 INT8 数据类型的线性 CHW，但 I/O 张量必须具有相同的数据类型。

bool PoolPlugin::supportsFormatCombination(
        int32_t pos, DynamicPluginTensorDesc const* inOut, int32_t nbInputs, int32_t nbOutputs)  noexcept override
{
    assert(nbInputs == 1 && nbOutputs == 1 && pos < nbInputs + nbOutputs);
    bool condition = inOut[pos].desc.format == PluginFormat::kLINEAR;
    condition &= (inOut[pos].desc.type == DataType::kFLOAT ||
                inOut[pos].desc.type == DataType::kHALF ||
            inOut[pos].desc.type == DataType::kBF16 ||
                inOut[pos].desc.type == DataType::kFP8 ||
                inOut[pos].desc.type == DataType::kINT8);
    condition &= inOut[pos].desc.type == inOut[0].desc.type;
    return condition;
}

重要提示

  • 如果 INT8 校准必须与具有 INT8 I/O 插件的网络一起使用，则插件必须支持 FP32 I/O，因为 TensorRT 使用 FP32 来校准图。

  • 如果不支持 FP32 I/O 变体或未使用 INT8 校准，则必须显式设置所有必需的 INT8 I/O 张量比例。

  • 校准无法确定插件内部张量的动态范围。在量化数据上运行的插件必须计算其内部张量的动态范围。

  • 插件可以设计为接受 FP8 和 INT8 I/O 类型，但请注意，在 TensorRT 9.0 中，构建器不允许混合 INT8 和 FP8 的网络。

TensorRT 通过 configurePlugin 或 onShapeChange 传递的信息可用于获取有关池化参数以及输入和输出比例的信息。这些信息可以存储为成员变量、序列化，然后在推理期间反序列化以供使用。

int32_t PoolPlugin::configurePlugin(DynamicPluginTensorDesc const* in, int32_t nbInputs, DynamicPluginTensorDesc const* out, int32_t nbOutputs) noexcept override
{
    ...
    mPoolingParams.mC = in.desc.d[1];
    mPoolingParams.mH = in.desc.d[2];
    mPoolingParams.mW = in.desc.d[3];
    mPoolingParams.mP = out.desc.d[2];
    mPoolingParams.mQ = ou.desc.d[3];
    mInHostScale = in[0].desc.scale >= 0.0F ? in[0].desc.scale : -1.0F;
    mOutHostScale = out[0].desc.scale >= 0.0F ? out[0].desc.scale : -1.0F;
}

已从 PluginTensorDesc::scale 获取每个张量的 INT8 I/O 比例。

使用 Python API 添加自定义层 (TensorRT >= 10.6)#

对于大多数用例，建议使用基于装饰器的方法定义 Python 插件（从 TensorRT 10.6 开始可用）。有关描述不同用例和最佳实践的手册，请参阅 TensorRT Python API 文档中的使用 TensorRT Python 插件编写自定义算子。

使用 Python API 添加自定义层（高级/TensorRT <= 10.5）#

使用基于类的方法的 Python（这也是 TensorRT <= 10.5 唯一支持的方法）。与基于装饰器的 Python 插件（在前一节中描述）相比，基于类的插件具有以下优点：

  • 状态性：基于类的插件具有状态（例如，已配置/未配置），并且 TensorRT 可以更精细地查询不同的插件属性和行为。

  • 形状张量输入支持。

  • 只有通过自定义插件创建器定义才能精细控制 TensorRT 在引擎反序列化期间创建的插件实例，而自定义插件创建器定义仅在使用基于类的方法时可用。

  • 插件属性的手动序列化和反序列化。

  • 能够预先请求设备内存暂存空间（除了输入/输出缓冲区之外的工作区），以避免执行时设备内存分配。

这些优势通常以增加实现复杂性和代码膨胀为代价，这可能导致更多错误。因此，在 Python 中考虑基于类的插件实现之前，建议进行权衡分析。

在 Python 中实现基于类的插件类似于 C++，都需要实现 IPluginV3 和 IPluginCreatorV3One。此外，Python 中的接口方法与其 C++ 对应方法具有大致相似的 API；大多数差异都很小且不言自明。

以下列表包含一些选定的更改。后续小节将更详细地描述所涉及的差异。

  • 以下插件 API 已被省略，取而代之的是读取/写入适当命名的属性。

类

方法

替换为属性

IPluginV3OneCore

getPluginName()

plugin_name[str]

IPluginV3OneCore

getPluginNamespace()

plugin_namespace [str]

IPluginV3OneCore

getPluginVersion()

plugin_version [str]

IPluginV3OneBuild

getNbOutputs()

num_outputs [int]

IPluginV3OneBuild

getTimingCacheID()

timing_cache_id [str]

IPluginV3OneBuild

getMetadataString()

metadata_string [str]

IPluginV3OneBuild

getFormatCombinationLimit()

format_combination_limit [int]

IPluginCreatorV3One

getPluginNamespace()

plugin_namespace [str]

IPluginCreatorV3One

getFieldNames()

field_names [PluginFieldCollection]

IPluginCreatorV3One

getPluginName()

name [str]

IPluginCreatorV3One

getPluginVersion()

plugin_version [str]

  • 一些方法具有默认实现；这些方法可以不实现，并且将采用以下概述的默认行为

    class trt.IPluginV3:
    def destroy(self):
        pass

class trt.IPluginV3OneBuild:
    def get_valid_tactics(self):
        return []

    def get_workspace_size(self, input_desc, output_desc):
        return 0

  • 在 IPluginV3OneBuild 和 IPluginV3OneRuntime 中必须返回整数状态代码的方法应在 Python 中引发异常。例如：

1int32_t configurePlugin(DynamicPluginTensorDesc const* in, int32_t nbInputs, DynamicPluginTensorDesc const* out, int32_t nbOutputs)

1configure_plugin(self: trt.IPluginV3OneBuild, in: List[trt.DynamicPluginTensorDesc], out: List[trt.DynamicPluginTensorDesc]) -> None

例如，如果输入具有非法值，则可以在 enqueue 期间引发 ValueError。

  • Python API IPluginV3.destroy() 在 C++ API 中没有直接等效项。Python 插件应在 IPluginV3.destroy() 方法中执行在 IPluginV3 C++ 析构函数中执行的任何功能。

有关演示 Python 插件的完整示例，请参阅 python_plugin 示例。

Python 插件的注册#

Python 插件必须通过 IPluginRegistry.register_creator() API 动态注册。没有类似于 REGISTER_TENSORT_PLUGIN 的模拟可用于静态注册。

构建和运行包含 Python 插件的 TensorRT 引擎#

可以使用基于 Python 的插件来构建 TensorRT 引擎。但是，目前无法在 Python 外部运行此类引擎，因为插件必须在引擎反序列化的范围内可用。例如，您不能直接使用像 trtexec 这样的工具。

实现 Python 插件的 enqueue#

C++ 和 Python 中 IPluginV3OneRuntime::enqueue() 的 API 如下：

1int32_t enqueue(PluginTensorDesc const *inputDesc, PluginTensorDesc const *outputDesc, void const *const *inputs, void *const *outputs, void *workspace, cudaStream_t stream)

1enqueue(self: trt.IPluginV3OneRuntime, input_desc: List[trt.PluginTensorDesc], output_desc: List[trt.PluginTensorDesc], inputs: List[int], outputs: List[int], workspace: int, stream: int) -> None

在此，inputs、outputs 和 workspace 作为各自设备指针的 intptr_t 强制类型转换传入。同样，stream 是指向 CUDA 流句柄的指针的 intptr_t 强制类型转换。在 Python 中，如何从这些缓冲区读取和写入具有灵活性，这可以根据特定的用例来实现。例如，使用 CUDA Python，这非常简单，因为 cuda.cuLaunchKernel 接受表示包装在 NumPy 数组中的指针的 int。

d_input = np.array([inputs[0]], dtype=np.uint64)
d_output = np.array([outputs[0]], dtype=np.uint64)
stream_ptr = np.array([stream], dtype=np.uint64)
args = [d_input,  d_output]
kernel_args = np.array([arg.ctypes.data for arg in args], dtype=np.uint64)
…
checkCudaErrors(cuda.cuLaunchKernel(_float_kernel,
                                        num_blocks, 1, 1,
                                        block_size, 1, 1,
                                        0,
                                        stream_ptr,
                                        kernel_args , 0))

将 enqueue 中的设备缓冲区/CUDA 流指针转换为其他框架#

可以使用 CuPy 的 UnownedMemory 类在设备缓冲区之上构建 CuPy 数组。

def enqueue(self, input_desc, output_desc, inputs, outputs, workspace, stream):
...
inp_dtype = trt.nptype(input_desc[0].type)
inp_mem = cp.cuda.UnownedMemory(
    inputs[0], volume(input_desc[0].dims) * cp.dtype(inp_dtype).itemsize, self
)
out_mem = cp.cuda.UnownedMemory(
    outputs[0],
    volume(output_desc[0].dims) * cp.dtype(inp_dtype).itemsize,
    self,
)

inp_ptr = cp.cuda.MemoryPointer(inp_mem, 0)
out_ptr = cp.cuda.MemoryPointer(out_mem, 0)

inp = cp.ndarray((volume(input_desc[0].dims)), dtype=inp_dtype, memptr=inp_ptr)
out = cp.ndarray((volume(output_desc[0].dims)), dtype=inp_dtype, memptr=out_ptr)

如果需要，然后可以使用 torch.as_tensor() 来构建 Torch 数组。

# inp_d = cp.ndarray(tuple(input_desc[0].dims), dtype=inp_dtype, memptr=inp_ptr)
inp_t = torch.as_tensor(inp_d, device='cuda')

类似地，可以通过 CuPy 的 ExternalStream 类从传入的流指针构造 CuPy 流句柄。

cuda_stream = cp.cuda.ExternalStream(stream)

自动向下转型#

TensorRT Python 绑定将为在 Python 中编写的实现 IPluginCreatorV3One 或 IPluginResource 等接口的自定义类型执行自动向下转型。例如，以 IPluginRegistry 中的以下方法为例：

get_creator(self: trt.IPluginRegistry, name: string, version: string,
namespace: string = “”) -> trt.IPluginCreatorInterface

返回类型指示为 IPluginCreatorInterface。但是，在实践中，如果您要编写一个实现 IPluginCreatorV3One 的类 MyPluginCreator（它又实现了 IPluginCreatorInterface），则 get_creator 方法将返回自动向下转型类型 MyPluginCreator。

这扩展到 trt.IPluginRegistry.all_creators，它是一个 List[trt.IPluginCreatorInterface]。如果您注册了一个类型为 MyPluginCreator 的插件创建器和另一个类型为 MyOtherPluginCreator 的插件创建器，则这两个插件创建器都将以列表中的相应类型存在。

示例：使用 Python 向 TensorRT 网络添加自定义层#

使用插件节点，可以将自定义层添加到 Python 中的任何 TensorRT 网络。Python API 有一个名为 add_plugin_v3 的函数，该函数允许向网络添加插件节点。以下示例对此进行了说明。它创建了一个简单的 TensorRT 网络，并通过查找 TensorRT 插件注册表添加了一个假设的插件节点。

import tensorrt as trt
import numpy as np

TRT_LOGGER = trt.Logger()

trt.init_libnvinfer_plugins(TRT_LOGGER, '')
def get_trt_plugin(plugin_name, plugin_version, plugin_namespace):
    plugin = None
    plugin_creator = trt.get_plugin_registry().get_creator(plugin_name, plugin_version, plugin_namespace)
    # trt will automatically downcast to IPluginCreator or IPluginCreatorInterface
    # Can inspect plugin_creator.interface_info to make sure
    if plugin_creator is not None:
        lrelu_slope_field = trt.PluginField("epsilon", np.array([0.00000001], dtype=np.float32), trt.PluginFieldType.FLOAT32)
        field_collection = trt.PluginFieldCollection([lrelu_slope_field])
        plugin = plugin_creator.create_plugin(name=plugin_name, field_collection=field_collection, phase=trt.TensorRTPhase.BUILD)
    return plugin

def main():
    builder = trt.Builder(TRT_LOGGER)
    network = builder.create_network()
    config = builder.create_builder_config()
    config.max_workspace_size = 2**20
    input_layer = network.add_input(name="input_layer", dtype=trt.float32, shape=(1, 1))
    plugin = network.add_plugin_v3(inputs=[input_layer], shape_inputs=[], plugin=get_trt_plugin("MY_PLUGIN", "1", ""))
    plugin.get_output(0).name = "outputs"
    network.mark_output(plugin.get_output(0))

启用时间缓存和使用自定义策略#

IPluginV3 提供了对自定义层的分析的更多控制，这在 V2 插件和更早版本中是不可用的。其中一项功能是启用时间缓存。如果 TensorRT 网络包含同一插件的多个实例，配置相同（例如，相同的插件属性值）并处理相同的输入输出形状和类型，则仅对一个实例进行时间测量（测量延迟）、缓存延迟并跳过其余实例的计时是有意义的。这可以大大节省引擎构建时间。

IPluginV3 插件的时间缓存是一项选择加入功能；要选择加入，插件必须声明一个非空的时间缓存 ID。

1char const* FooPlugin::getTimingCacheID() noexcept override
2{
3    // return nullptr to disable timing caching (default behavior)
4    // return non-null string to enable timing caching
5}

1def FooPlugin(trt.IPluginV3, trt.IPluginV3OneBuild, ...):
2    def __init__(self):
3        # set to None to disable timing caching
4        self.timing_cache_id = value

请注意以下关于时间缓存 ID 的事项：

  • 用户提供的时间缓存 ID 应被视为较大缓存 ID 的后缀；TensorRT 通过考虑插件的输入/输出形状和格式信息自动形成前缀。通常，用户提供的时间缓存 ID 可以包含插件属性及其值。

  • 它必须反映插件的创建状态，并且在创建后不会演变。

对于 V2 插件，TensorRT 仅对插件声明支持的任何（多个）类型/格式组合进行计时。使用 IPluginV3，插件还可以确保自定义策略被计时，并且 TensorRT 使用最快的策略。例如，插件可能具有两个内核之一来计算输出，并且可能无法预测哪个内核在特定平台以及特定输入/输出形状和格式下最快。可以要求 TensorRT 对每个策略的每种格式组合计时插件，找出最快的此类配置，并在推理期间使用它。

注意

  • 如果 TensorRT 仅支持一种类型/格式组合，并且要么不使用自定义策略，要么仅声明一种策略，则 TensorRT 可能会选择不对插件进行计时。

  • 对于 IPluginV3OneBuild，TensorRT 最多计时 getFormatCombinationLimit() 种类型/格式组合用于每个策略；根据需要覆盖此方法以增加/减少此限制。

要开始使用，请向 TensorRT 声明自定义策略：

 1int32_t FooPlugin::getNbTactics() noexcept override
 2{
 3    return 2; // return 0 to disable custom tactics (default behavior)
 4}
 5
 6int32_t FooPlugin::getValidTactics(int32_t* tactics, int32_t nbTactics) noexcept override
 7{
 8    tactics[0] = 1;
 9    tactics[1] = 2;
10    return 0;
11}

1def get_valid_tactics(self):
2    return [1, 2] # return empty vector to disable custom tactics (default behavior)

任何严格正整数都可以用作自定义策略值（TensorRT 保留 0 作为默认策略）。

当插件被计时时，保证在调用 getValidTactics() 之前使用当前输入/输出格式组合调用 configurePlugin()。因此，可以为每个输入/输出格式组合声明一组不同的策略。例如，对于支持 FP32 和 FP16 的插件，策略 1 可能仅限于 FP16，同时支持 FP32 的策略 1 和 2。

在引擎构建期间，当自动调整插件时，TensorRT 将通过调用 IPluginV3OneRuntime::setTactic (C++，Python) 传达后续 enqueue() 的策略。当引擎反序列化时，TensorRT 将在创建插件后调用 setTactic，以传达为插件选择的最佳策略。即使未使用自定义策略，也会使用默认策略值 0 调用 setTactic。

在插件之间共享自定义资源#

从 TensorRT 10.0 开始，键值存储与插件注册表关联。此存储可以存储用户实现的 IPluginResource (C++，Python) 对象，并使用字符串键进行索引。此功能可用于在不同插件之间共享状态或某些资源。请注意，它不与 IPluginV3（甚至插件接口）绑定。

让我们探索一个示例。

示例：在通过网络下载的不同插件之间共享权重#

假设多个插件需要访问相同的权重 W。由于许可限制，您可能希望在引擎运行时下载这些权重。但是，由于 W 的尺寸很大，因此也希望只下载一个副本，该副本在所有需要访问的插件之间共享。

  1. 实现 SharedWeights 类，该类实现 IPluginResource。

  2. 每个需要访问权重的插件都通过调用 IPluginRegistry::acquirePluginResource(...) (C++，Python) 请求已初始化（已下载）的 SharedWeights 实例。

1IPluginResource* acquirePluginResource(char const* key, IPluginResource* resource)

1acquire_plugin_resource(self: trt.IPluginRegistry, key: str, resource: trt.IPluginResource) -> trt.IPluginResource

第一次针对特定 key 调用 acquirePluginResource 时，TensorRT 会注册提供的插件 resource 的克隆，而不是作为 resource 传递的对象。注册的对象通过调用 resource->clone() 获得。因此，最佳实践是仅初始化克隆 - 在这种情况下，权重下载可以在 IPluginResource::clone() 中完成。

  1. 每个插件完成权重使用后，它可以调用 IPluginRegistry::releasePluginResource() 以指示它不再希望使用它们。

1int32_t releasePluginResource(char const* key)

1release_plugin_resource(self: trt.IPluginRegistry, key: str) -> None

TensorRT 对针对特定键进行的 acquirePluginResource 和 releasePluginResource 调用执行引用计数，并且当引用计数达到零时，将调用 IPluginResource::release()。在此示例中，可以利用此功能来释放权重使用的内存，前提是所有插件都已完成使用。

  1. 最后，SharedWeights 类可以实现如下：

    class SharedWeights : public IPluginResource
{
public:
    SharedWeights(bool init = false)
    {
        if(init)
        {
            PLUGIN_CHECK(cudaMalloc((void**) &cloned->mWeights, ...));
        }
    }

    int32_t release() noexcept override
    {
        TRY
        {
            if (mWeights != nullptr)
            {
                PLUGIN_CHECK(cudaFree(mWeights));
                mWeights = nullptr;
            }
        }
        CATCH
        {
            return -1;
        }
        return 0;
    }

    IPluginResource* clone() noexcept override
    {
        TRY
        {
            auto cloned = std::make_unique<SharedWeights>(/* init */ true);
            //
            // Download the weights
            //
            // Copy to device memory
            PLUGIN_CHECK(cudaMemcpy(cloned->mWeights, ...));
        }
        CATCH
        {
            return nullptr;
        }
        return cloned.release();
    }

    ~SharedWeights() override
    {
        if(mWeights)
        {
            release();
        }
    }

    float* mWeights{nullptr};
};

假设 FooPlugin 需要访问权重。它可以在准备好进行推理时请求权重。这可以在 IPluginV3OneRuntime::onShapeChange 中完成，该方法将在构建和运行时阶段至少调用一次即将 enqueue() 的插件。

int32_t onShapeChange(
    PluginTensorDesc const* in, int32_t nbInputs, PluginTensorDesc const* out, int32_t nbOutputs) noexcept override
{
    SharedWeights w{};
    mW = static_cast<SharedWeights*>(getPluginRegistry()->acquirePluginResource("W", &w))->mWeights;
    return 0;
}

然后可以在后续的 enqueue() 中使用获得的权重 (mW)。总结一下，插件可以在其析构函数中发出释放意图信号（请注意，在 IPluginV3 中没有类似于 IPluginV2DynamicExt::terminate() 的单独的释放资源例程）。

~FooPlugin() override
{
    TRY
    {
        PLUGIN_CHECK(getPluginRegistry()->releasePluginResource("W"));
    }
    CATCH
    {
        // Error handling
    }
}

所有需要访问权重的插件都可以使用上面的相同代码。引用计数机制将确保权重的可用性和正确释放。

使用解析器导入模型时使用自定义层#

ONNX 解析器自动尝试将无法识别的节点作为插件导入。如果在插件注册表中找到与节点具有相同 op_type 的插件，则解析器会将节点的属性作为插件字段参数转发到插件创建器，以创建插件。默认情况下，解析器使用 "1" 作为插件版本，使用 """ 作为插件命名空间。可以通过在相应的 ONNX 节点中设置 plugin_version 和 plugin_namespace 字符串属性来覆盖此行为。

有时，您可以在将 ONNX 图形导入 TensorRT 之前对其进行修改。例如，用插件节点替换一组 ops。为此，您可以使用 ONNX GraphSurgeon 实用程序。有关如何使用 ONNX-GraphSurgeon 替换子图的详细信息，请参阅此示例。

有关更多示例，请参阅 onnx_packnet 示例。

插件 API 描述#

所有新插件都应从 IPluginCreatorV3One 和 IPluginV3 类派生。此外，新插件还应在插件注册表中注册，可以通过使用 IPluginRegistry::registerCreator() 动态注册，也可以使用 REGISTER_TENSORRT_PLUGIN(...) 宏静态注册。自定义插件库还可以考虑实现等效于 initLibNvInferPlugins() 的 init 函数，以执行批量注册。

注意

汽车安全用户必须使用 REGISTER_SAFE_TENSORRT_PLUGIN(...) 宏而不是 REGISTER_TENSORRT_PLUGIN(...)。

IPluginV3 API 描述#

以下部分描述了 IPluginV3 以及扩展的 IPluginV3OneCore、IPluginV3OneBuild 或 IPluginV3OneBuildV2 和 IPluginV3OneRuntime 的功能。

由于 IPluginV3 对象由不同的功能组成，因此可以在其生命周期的任何时候调用 IPluginV3::getCapabilityInterface。为构建阶段添加的 IPluginV3 对象必须为所有功能类型返回有效的功能接口：核心、构建和运行时。对于为运行时阶段添加的对象，可以省略构建功能。

有一些方法用于请求有关插件的识别信息。它们也可以在插件生命周期的任何阶段调用。

  • IPluginV3OneCore::getPluginName：用于查询插件的名称

  • IPluginV3OneCore::getPluginVersion：用于查询插件的版本

  • IPluginV3OneCore::getPluginNamespace：用于查询插件的命名空间

  • IPluginV3OneBuild::getMetadataString：用于查询与插件关联的任何元数据的字符串表示形式，例如其属性的值。

为了将插件层连接到相邻层并设置输入和输出数据结构，构建器通过调用以下插件方法来检查输出数量及其形状：

  • IPluginV3OneBuild::getNbOutputs：用于指定输出张量的数量。

  • IPluginV3OneBuild::getOutputShapes：此函数将输出形状指定为输入形状或常量的函数。数据相关的形状具有指定的上限和最佳调整值除外。

  • IPluginV3OneBuild::supportsFormatCombination：用于检查插件是否支持给定的数据类型和格式组合。

  • IPluginV3OneBuild::getOutputDataType：此函数检索输出张量的数据类型。返回的数据类型必须是插件支持的格式。

如果使用 IPluginV3OneBuildV2 构建功能，则插件还可以向 TensorRT 传达某些输入输出对是别名（共享相同的数据缓冲区）。TensorRT 将查询 IPluginV3OneBuildV2::getAliasedInput 以确定任何此类别名行为。要使用此功能，必须启用 PreviewFeature::kALIASED_PLUGIN_IO_10_03。

插件层可以支持以下数据格式：

  • LINEAR 单精度 (FP32)、半精度 (FP16)、脑浮点 (BF16)、8 位浮点 E4M3 (FP8)、整数 (INT8) 和整数 (INT32) 张量

  • CHW32 单精度 (FP32) 和整数 (INT8) 张量

  • CHW2、HWC8、HWC16 和 DHWC8 半精度 (FP16) 张量

  • CHW4 半精度 (FP16) 和整数 (INT8) 张量

  • HWC8、HWC4、NDHWC8、NC2HW 脑浮点 (BF16) 张量

PluginFormat 计算格式。

不就地计算所有数据并且除了输入和输出张量之外还需要内存空间的插件可以使用 IPluginV3OneBuild::getWorkspaceSize 方法指定其他内存需求，构建器调用该方法来确定和预分配暂存空间。

层在构建时配置、执行和销毁，以发现最佳配置。在为插件选择最佳配置后，所选策略和具体形状/格式信息（数据相关的维度除外）会在推理期间传达给插件。它在推理应用程序的生命周期内根据需要执行多次，并在引擎销毁时最终销毁。

构建器使用以下插件方法控制这些步骤和运行时。推理期间也调用的方法用 (*) 表示 - 所有其他方法仅由构建器调用。

  • IPluginV3OneBuild::attachToContext*：此函数请求将插件克隆附加到 ExecutionContext，从而允许插件访问任何特定于上下文的资源。

  • IPluginV3OneBuild::getTimingCacheId：此函数查询 TensorRT 可能使用的任何时间缓存 ID。如果提供，则启用时间缓存（默认情况下禁用）。

  • IPluginV3OneBuild::getNbTactics：用于查询插件选择使用的自定义策略的数量。

  • IPluginV3OneBuild::getValidTactics：此函数查询插件可能使用的任何自定义策略。将为每个策略分析插件的性能，最多为 IPluginV3OneBuild::getFormatCombinationLimit() 指示的最大值。

  • IPluginV3OneBuild::getFormatCombinationLimit：此函数查询每个策略可能计时的最大格式组合数（如果未为默认策略声明自定义策略，则为 0）。

  • IPluginV3OneRuntime::setTactic*：传达要在后续 enqueue() 期间使用的策略。如果未声明自定义策略，则这始终为 0。

  • IPluginV3OneBuild::configurePlugin：传达输入和输出的数量及其形状、数据类型和格式。min、opt 和 max 以及每个输入或输出的 DynamicPluginTensorDesc 对应于插件当前分析的优化配置文件的 kMIN、kOPT 和 kMAX 值。desc.dims 字段对应于在网络创建时指定的插件输入的维度。在此阶段，desc.dims 字段中可能存在通配符维度。

    此时，插件可以设置其内部状态，并为给定的配置选择最合适的算法和数据结构。

  • IPluginV3OneRuntime::onShapeChange*：传达输入和输出的数量及其形状、数据类型和格式。维度是具体的，除非存在数据相关的维度，通配符将指示这些维度。

  • IPluginV3OneRuntime::enqueue*：封装插件的实际算法和内核调用，并提供指向输入、输出和暂存空间以及 CUDA 流的指针，CUDA 流将用于内核执行。

  • IPluginV3::clone：每次创建包含此插件层的新构建器、网络或引擎时都会调用此方法。它必须返回具有正确参数的新插件对象。

在构建器完成分析后，在引擎序列化之前，会调用 IPluginV3OneRuntime::getFieldsToSerialize 以查询必须序列化到引擎中的任何插件字段。这些字段应是插件在引擎反序列化后在推理阶段正常运行所需的数据。

IPluginCreatorV3One API 描述#

IPluginCreatorV3One 类中的以下方法用于从插件注册表中查找和创建适当的插件：

  • getPluginName：此方法返回插件名称，应与 IPluginV3OneCore::getPluginName 的返回值匹配。

  • getPluginVersion：返回插件版本。对于所有内部 TensorRT 插件，此值默认为 1。

  • getPluginNamespace：返回插件命名空间。默认值可以为 ""。

  • getFieldNames：要成功创建插件，您必须知道所有插件的字段参数。此方法返回 PluginFieldCollection 结构，其中填充了 PluginField 条目以反映字段名称和 PluginFieldType（数据应指向 nullptr）。

  • createPlugin：此方法创建一个插件，并传递 PluginFieldCollection 和 TensorRTPhase 参数。

在引擎反序列化期间，TensorRT 调用此方法，并将 TensorRTPhase 参数设置为 TensorRTPhase::kRUNTIME，并将 PluginFieldCollection 填充为与 IPluginV3OneRuntime::getFieldsToSerialize() 返回的 PluginFields 相同的值。在这种情况下，TensorRT 将获取由 createPlugin 返回的插件对象的所有权。

您还可以调用 createPlugin 来生成插件对象，以添加到 TensorRT 网络中。在这种情况下，建议将 phase 参数设置为 TensorRTPhase::kBUILD。使用 PluginFieldCollection 传递的数据应由调用者在程序销毁之前分配和释放。createPlugin 函数返回的插件对象的所有权将传递给调用者，并且必须由调用者销毁。

将 V2 插件迁移到 IPluginV3#

自 TensorRT 8.5 起，IPluginV2 和 IPluginV2Ext 已被弃用，而 IPluginV2IOExt 和 IPluginV2DynamicExt 在 TensorRT 10.0 中已被弃用。因此，新插件应以 IPluginV3 为目标，旧插件应进行重构。

将 IPluginV2DynamicExt 插件迁移到 IPluginV3 时，请记住以下关键点

  • 与插件关联的插件创建器必须迁移到 IPluginCreatorV3One，即 IPluginV3 的工厂类（IPluginCreator 是 IPluginV2 派生类的工厂类）。这仅包括迁移 IPluginCreator::deserializePlugin。有关更多信息，请参阅 插件序列化和反序列化 部分。

  • 在 IPluginV3 中，没有与 IPluginV2::initialize()、IPluginV2::terminate() 和 IPluginV2::destroy() 等效的方法。有关更多信息，请参阅 插件初始化和终止 部分。

  • 在 IPluginV3 中，没有与 IPluginV2Ext::detachFromContext() 等效的方法。有关更多信息，请参阅 访问 TensorRT 提供的上下文特定资源 部分。

  • IPluginV3OneRuntime::attachToContext() 在参数和行为方面与 IPluginV2Ext::attachToContext() 显著不同。有关更多信息，请参阅 访问 TensorRT 提供的上下文特定资源 部分。

  • 在 IPluginV3 中，插件序列化通过 PluginFieldCollection 进行，该集合通过 IPluginV3OneRuntime::getFieldsToSerialize() 传递给 TensorRT，反序列化通过相同的 PluginFieldCollection 进行，该集合由 TensorRT 传递回 IPluginCreatorV3One::createPlugin(...)。有关更多信息，请参阅 插件序列化和反序列化 部分。

  • IPluginV3 中 IPluginV2DynamicExt 中 void 返回方法的等效方法将期望整数状态代码作为返回值（例如，configurePlugin）。

  • supportsFormatCombination 和 getWorkspaceSize 获取动态张量描述符 (DynamicPluginTensorDesc) 而不是静态描述符 (PluginTensorDesc)。

  • IPluginV2DynamicExt::getOutputDimensions() 变为 IPluginV3OneBuild::getOutputShapes()，并更改为输出参数签名而不是返回值。它还从按输出索引查询转变为一次性查询。类似的转换也适用于从 IPluginV2Ext::getOutputDataType 到 IPluginV3OneBuild::getOutputDataTypes。

插件初始化和终止#

IPluginV2 提供了几个用于插件初始化和终止的 API：即 IPluginV2::initialize()、IPluginV2::terminate() 和 IPluginV2::destroy()。在 IPluginV3 中，插件预计在初始化状态下构建；如果您的 V2 插件在 initialize 中有任何延迟初始化，则可以将其推迟到 onShapeChange 或 configurePlugin。 IPluginV2::terminate() 或 IPluginV2::destroy() 中的任何资源释放或终止逻辑都可以移动到类析构函数中。Python API 中存在例外；IPluginV3.destroy() 作为 C++ 风格析构函数的替代方案提供。

访问 TensorRT 提供的上下文特定资源#

IPluginV2Ext::attachToContext() 为插件提供了对上下文特定资源的访问权限，即 GPU 分配器以及 cuDNN 和 cuBLAS 句柄。IPluginV3OneRuntime::attachToContext() 旨在为插件提供类似的服务，但它提供的是 IPluginResourceContext，后者又公开了插件可能请求的资源。

与 IPluginV2Ext::attachToContext() 不同的是，IPluginResourceContext 不再提供 cuDNN 和 cuBLAS 句柄；任何依赖于这些句柄的插件都应迁移以初始化自己的 cuDNN 和 cuBLAS 资源。如果首选在插件之间共享 cuDNN/cuBLAS 资源，则可以利用 IPluginResource 和插件注册表的键值存储提供的功能来实现此目的。有关更多信息，请参阅 在插件之间共享自定义资源 部分。

IPluginV3OneRuntime::attachToContext(...) 是一种克隆和附加操作。它被要求克隆整个 IPluginV3 对象——而不仅仅是运行时功能。因此，如果作为单独的类实现，运行时功能对象可能需要保持对其所属的 IPluginV3 对象的引用。

通过 IPluginResourceContext 获取的任何上下文特定资源都可以使用，直到插件被销毁。因此，在 IPluginV2Ext::detachFromContext() 中实现的任何终止逻辑都可以移动到插件析构函数中。

插件序列化和反序列化#

对于 V2 插件，序列化和反序列化由 IPluginV2::serialize、IPluginV2::getSerializationSize 和 IPluginCreator::deserializePlugin 的实现决定；IPluginV3OneRuntime::getFieldsToSerialize 和 IPluginCreatorV3One::createPlugin 已替换这些。请注意，工作流程已从写入/读取原始缓冲区转变为构造和解析 PluginFieldCollection。

TensorRT 处理在 PluginFieldType 中定义的类型的序列化。自定义类型可以序列化为 PluginFieldType::kUNKNOWN。例如

struct DummyStruct
{
    int32_t a;
    float b;
};

DummyPlugin()
{
    // std::vector<nvinfer1::PluginField> mDataToSerialize;
    // int32_t mIntValue;
    // std::vector<float> mFloatVector;
    // DummyStruct mDummyStruct;
    mDataToSerialize.clear();
    mDataToSerialize.emplace_back(PluginField("intScalar", &mIntValue, PluginFieldType::kINT32, 1));
    mDataToSerialize.emplace_back(PluginField("floatVector", mFloatVector.data(), PluginFieldType::kFLOAT32, mFloatVector.size()));
    mDataToSerialize.emplace_back(PluginField("dummyStruct", &mDummyStruct, PluginFieldType::kUNKNOWN, sizeof(DummyStruct)));
    mFCToSerialize.nbFields = mDataToSerialize.size();
    mFCToSerialize.fields = mDataToSerialize.data();
}

nvinfer1::PluginFieldCollection const* DummyPlugin::getFieldsToSerialize() noexcept override
{
    return &mFCToSerialize;
}

将旧版 V2 插件迁移到 IPluginV3#

如果从 IPluginV2 或 IPluginV2Ext 迁移到 IPluginV3，则最好先迁移到 IPluginV2DynamicExt，然后按照上述指南迁移到 IPluginV3。IPluginV2DynamicExt 中的新功能如下

virtual DimsExprs getOutputDimensions(int outputIndex, const DimsExprs* inputs, int nbInputs, IExprBuilder& exprBuilder) = 0;

virtual bool supportsFormatCombination(int pos, const PluginTensorDesc* inOut, int nbInputs, int nbOutputs) = 0;

virtual void configurePlugin(const DynamicPluginTensorDesc* in, int nbInputs, const DynamicPluginTensorDesc* out, int nbOutputs) = 0;

virtual size_t getWorkspaceSize(const PluginTensorDesc* inputs, int nbInputs, const PluginTensorDesc* outputs, int nbOutputs) const = 0;

virtual int enqueue(const PluginTensorDesc* inputDesc, const PluginTensorDesc* outputDesc, const void* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) = 0;

迁移到 IPluginV2DynamicExt 的指南是

  • getOutputDimensions 实现给定输入的输出张量维度的表达式。

  • supportsFormatCombination 检查插件是否支持指定 I/O 的格式和数据类型。

  • configurePlugin 模仿 IPluginV2Ext 中等效 configurePlugin 的行为，但接受张量描述符。

  • getWorkspaceSize 和 enqueue 模仿 IPluginV2Ext 中等效 API 的行为，但接受张量描述符。

插件编码指南#

内存分配

必须释放插件中分配的内存，以确保没有内存泄漏。如果在插件构造函数或后续阶段（如 onShapeChange）获取资源，则必须释放这些资源，可能在插件类析构函数中释放。

另一种选择是通过 getWorkspaceSize 请求所需的任何额外工作区内存，这将在 enqueue 期间可用。

添加检查以确保正确配置并验证输入

插件意外行为的常见来源是不正确的配置（例如，无效的插件属性）和无效的输入。因此，在插件开发的初始阶段，最好为插件预期无法工作的情况添加检查/断言。以下是可能添加检查的位置

  • createPlugin：插件属性检查

  • configurePlugin 或 onShapeChange：输入维度检查

  • enqueue：输入值检查

对于创建新插件对象的方法，在错误时返回 Null

诸如 createPlugin、clone 和 attachToContext 之类的方法可能需要创建并返回新的插件对象。在这些方法中，如果出现任何错误或检查失败，请确保返回空对象（C++ 中的 nullptr）。这确保在配置不正确时不会返回非空插件对象。

避免在 clone() 中分配设备内存

由于构建器多次调用 clone，因此设备内存分配可能非常昂贵。一种选择是在构造函数中执行持久内存分配，在插件准备就绪时（例如，在 configurePlugin 中）复制到设备，并在销毁期间释放。

序列化任意数据块和自定义类型

插件作者可以利用 PluginFieldType::kUNKNOWN 的 PluginField 来指示要序列化的任意数据块。在这种情况下，相应 PluginField 的 length 应为与 data 指向的缓冲区对应的字节数。可以通过这种方式实现非原始类型的序列化。

插件共享库#

TensorRT 包含可以静态加载到应用程序中的内置插件。

您可以使用 REGISTER_TENSORRT_PLUGIN 和 registerCreator 接口显式向 TensorRT 注册自定义插件（请参阅 使用 C++ 添加自定义层）。但是，您可能希望 TensorRT 管理插件库的注册，特别是将插件库与 plan 文件一起序列化，以便在创建引擎时自动加载它们。当您想要将插件包含在版本兼容的引擎中时，这尤其有用，这样您就不需要在构建引擎后管理它们。要利用这一点，您可以构建具有 TensorRT 识别的特定入口点的共享库。

生成插件共享库#

要为插件创建共享库，该库必须定义以下公共符号

extern "C" void setLoggerFinder(ILoggerFinder* finder);
extern "C" IPluginCreator* const* getCreators(int32_t& nbCreators) const;

上面的 extern "C" 仅用于防止名称修饰，这些方法应在 C++ 中实现。有关更多详细信息，请查阅编译器的 ABI 文档。

setLoggerFinder() 应在库中设置 ILoggerFinder 的全局指针，以便在插件代码中进行日志记录。getPluginCreators() 返回库包含的插件创建器列表。这些入口点的示例可以在 plugin/common/vfcCommon.h/cpp 中找到。

要将插件库与引擎 plan 一起序列化，请使用 BuilderConfig 中的 setPluginsToSerialize() 向 TensorRT 提供插件库路径。

您还可以在构建版本兼容的引擎时将插件打包到 plan 中。打包的插件将与引擎具有相同的生命周期，并且在运行引擎时将自动注册/注销。

使用插件共享库#

构建共享库后，您可以配置构建器以将其与引擎一起序列化。下次将引擎加载到 TensorRT 中时，将自动加载和注册序列化的插件库。

注意

IPluginRegistry loadLibrary() (C++, Python) 功能现在支持包含 V2 和 V3 插件创建器的插件共享库，通过 getCreators() 入口点。getPluginCreators() 入口点也有效，但已被弃用。TensorRT 首先检查 getCreators() 符号是否可用，如果不可用，则检查 getPluginCreators() 作为向后兼容的后备选项。然后，您可以查询此内容以枚举每个插件创建器，并使用 IPluginRegistry registerCreator() (C++, Python) 手动注册它。

在构建引擎之前，加载插件以供构建器使用

1for (size_t i = 0; i < nbPluginLibs; ++i)
2{
3    builder->getPluginRegistry().loadLibrary(pluginLibs[i]);
4}

1for plugin_lib in plugin_libs:
2    builder.get_plugin_registry().load_library(plugin_lib)

接下来，决定是否应将插件包含在引擎中或在外部发布。您可以按如下方式将插件与 plan 一起序列化

1IBuilderConfig *config = builder->createBuilderConfig();
2...
3config->setPluginsToSerialize(pluginLibs, nbPluginLibs);

1config = builder.create_builder_config()
2...
3config.plugins_to_serialize = plugin_libs

或者，您可以将插件保留在引擎外部。您需要在部署引擎时随引擎一起发布这些库，并在反序列化引擎之前在运行时显式加载它们

1// In this example, getExternalPluginLibs() is a user-implemented method which retrieves the list of libraries to use with the engine
2std::vector<std::string> pluginLibs = getExternalPluginLibs();
3for (auto const &pluginLib : pluginLibs)
4{
5    runtime->getPluginRegistry().loadLibrary(pluginLib.c_str())
6}

1# In this example, get_external_plugin_libs() is a user-implemented method which retrieves the list of libraries to use with the engine
2plugin_libs = get_external_plugin_libs()
3for plugin_lib in plugin_libs:
4    runtime.get_plugin_registry().load_library(plugin_lib)
本页内容