[TensorRT] 构建推理网络中（自定义层）Plugin Layer使用方法及数据处理格式

最新推荐文章于 2023-05-15 16:41:29 发布

xuke_2022

最新推荐文章于 2023-05-15 16:41:29 发布

阅读量1.1k

点赞数 1

分类专栏： AI 文章标签：机器学习

本文链接：https://blog.csdn.net/xuke_2018/article/details/80255926

版权

本文详细介绍了在TensorRT中构建推理网络时如何使用自定义层（Plugin Layer），包括nvinfer.h中的关键定义、数据维度处理、Plugin与PluginFactory的实现，以及如何利用Plugin构造和配置推理网络。在引擎序列化过程中，着重讲解了PluginLayer的输出维度计算、配置方法和序列化流程，为理解TensorRT中自定义操作的实现提供了清晰的指导。

摘要由CSDN通过智能技术生成

代码详见SampleCharRNN

1 nvinfer.h中的定义

1.1 处理数据维度的定义

/**
 * \class Dims
 * \brief structure to define the dimensions of a tensor
 *
 * \note: currently the following formats are supported for layer inputs and outputs:
 * * zero or more index dimensions followed by one channel and two spatial dimensions (e.g. CHW)
 * * one time series dimension followed by one index dimension followed by one channel dimension (i.e. TNC)
 */
class Dims
{
public:
    static const int MAX_DIMS = 8; //!< the maximum number of dimensions supported for a tensor
    int nbDims;                    //!< the number of dimensions
    int d[MAX_DIMS];               //!< the extent of each dimension
    DimensionType type[MAX_DIMS];  //!< the type of each dimension
};

1.2 Infer engine中加入PluginLayer