TensorRT文档翻译

最新推荐文章于 2025-04-28 20:17:20 发布

seekforcode

最新推荐文章于 2025-04-28 20:17:20 发布

阅读量2k

点赞数 1

CC 4.0 BY-SA版权

分类专栏： TensorRT 文章标签： TensorRT

本文链接：https://blog.youkuaiyun.com/seekforcode/article/details/79643642

本文档翻译自TensorRT 3.8章节，详细解释如何实现TensorRT的自定义层，包括创建网络、配置、运行时实现和序列化等步骤，并通过SamplePlugin示例介绍自定义层的完整流程。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

1 首先声明，我也是刚开始研究这个TensorRT的自定义层，之前试了很多方法，下了官方例程，看不懂ORZ。然后也去找了一些别的资料，也没能复现。特别是把NVIDIA的官方例子face_recongnition，在本机上编译失败（ubuntu16.04+gtx1060（pascal架构）），可能是因为我没有装OpenGL还有没有配置好tensorRT，但是在jetson tx1上面却是可以的0.0。

2 其次我翻译的是tensorRT3里面的TensorRT-3-User-Guide.pdf里面关于自定义层的说明，也就是3.8章节

正文

3.8. SamplePlugin - Implementing A Custom Layer
这个例子示范了怎么给tensorRT添加一个自定义层。它用直接调用cuda基础线性代数库来代替原来MNIST例子中的最后的全连接层

3.8.1关键概念
这里的关键概念为了解释以下两点：
1如何去创建自定义层
2如何去整合自定义层到NvCaffeParser然后把它输入到执行文件中

3.8.2实现plugin接口
tensorRT每个阶段（网络创建，建立，序列化，执行）需要在plugin API中实现的方法。下面回顾以下各个阶段的相关方法。

当创建网络的时候：
当网络构建的时候这些方法被调用，如果层的输出size或者任何随后的层需要通过一个ITensor::getDimensions()的调用。另外这个方法将被调用当建立者（builder）运行时。
‣ getNbOutputs()
‣ getOutputDimensions()

builder的话：

‣ configure()
‣ getWorkspaceSize()

在runtime：

‣ initialize() when an engine context is constructed
‣ enqueue() at inference time
‣ terminate() when an engine context is destroyed

For serialization（怎么翻译？？？序列化阶段？）:

‣ getSerializationSize()
‣ serialize()

3.8.3定义网络
当在定义网络的时候，TensorRT需要去知道哪一层是输出层
/************************************************************************************************************************************************************/
*sample中给的维度没有batch size，用了相同的方法ITensor::getDimensions()去返回维度。举个栗子，一个典型的3维的卷积，维度被提供以{C,H,W}的形式（通道，高，宽），
*然后返回值也应该是{C,H,W}的格式
/************************************************************************************************************************************************************/

下面的方法提供了层的输出
int getNbOutputs() const override
{
return 1;
}

getOutputDimensions函数有三个参数，输出指针，输入维度和输出的数量。后面两个参数早就被TensorRT内部计算好了，所以你要做的是计算输出的维度基于输入维度和指针数量。

Dims getOutputDimensions(int index, const Dims* inputDims, int nbInputDims) override//override是什么鬼0.0，我学的是假的c++吧
{