TensorRT文档翻译

seekforcode

墨衍会员 · AI 创作全网分发码龄12年

关注

原创最新推荐文章于 2026-07-27 11:55:00 发布 2.1k 阅读 AI 写同款· GEO 优化 ›

TensorRT从了入门到了解(官方文档)-学习笔记 TensorRT阅读研究阅读详情

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++吧
{
assert(index == 0 && nbInputDims == 1 && inputDims[0].nbDims== 3);
assert(mNbInputChannels == inputDims[0].d[0] *inputDims[0].d[1] *inputDims[0].d[2]);
return DimsCHW(mNbOutputChannels, 1, 1);
}

3.8.4.层构建（ Layer Configuration）

建立者（builder）调用网络的configuration()方法（ps：这貌似不是c++的说法吧，方法感觉应该是java的说法吧，叫成员函数是不是更好一点），去给一个选择基于他的输入的算法的机会。
在这个例子中，输入被检查是否是正确的格式。在更加复杂的例子中，你可能选择基于输入维度的卷积算法。

这个configure()方法是在build阶段唯一被调用的方法。因此，任何决定这里的东西也会在runtime阶段被需要去保存（stored）。作为plugin的成员变量，serialized（序列化？？串连？？）以及de-serialized亦然

3.8.5工作空间（Workspace）

TensorRT提供了工作空间为了在层执行期间的临时存储，为了最小化内存占用被各个层共享。TensorRT builder 调用getWorkspaceSize()为了决定workspace的需求（requirement）（毛估估判定需要的大小吧，类似于分配空间）。
在这个例子中，没有用到workspace。如果需要workspace，它将会被分配当一个IExecutionContext被创建，然后在runtime阶段转到enqueue()方法。

非此即彼（Alternatively），你也可以给层的执行分配显存在构造函数中，然后再析构函数中释放。然后只需要在这个回调函数（callback）中返回0。使用TensorRT的worksoace的优势是可以在执行的时候被其他plugins分享。

3.8.6 资源管理（Resource Management）
initialize()和 terminate()方法被builder调用当展示自动调整（performing auto-tuning），并且在runtime阶段当IExecutionContext被创建和销毁时。（我理解为两个阶段都会调用这两个函数）他们被用于分配和释放自愿在层执行时。
在这个例子中，为了cuDNN和cuBLAS，还有一些cuDNN张量（tensor）描述符为了bias额外操作的的句柄（handles）。

在下面的例子中，为了cuDNN和cuBLAS，还有一些为了bias额外操作的cuDNN张量（tensor）描述符创建了句柄（handles)。（为cudnn，cublas，bias额外操作提供的句柄）（ps，blablablabla没理解）

int initialize() override
{
CHECK(cudnnCreate(&mCudnn));
CHECK(cublasCreate(&mCublas));
CHECK(cudnnCreateTensorDescriptor(&mSrcDescriptor));
CHECK(cudnnCreateTensorDescriptor(&mDstDescriptor));
return 0;
}
virtual void terminate() override
{
CHECK(cublasDestroy(mCublas));
CHECK(cudnnDestroy(mCudnn));
}

3.8.7.runtime阶段实现（Runtime Implementation）

enqueue()方法会被使用去执行层的runtim阶段的实现
batch size 会被传到enqueue（）以最多为最大指定batch size在build阶段，尽管他可以更小

/************************************************************************************************************************************************************/
*除了batch size，尺寸信息(dimensional information)不会被传到enqueue(),所以，其他的尺寸信息在runtime阶段被需要，举个栗子，输入和输出的通道数目，需要作为层数据
*的一部分被串连（serialized）
/************************************************************************************************************************************************************/

virtual int enqueue(int batchSize, const void*const * inputs,void**outputs, void* workspace, cudaStream_t stream) override
{
int nbOutputChannels = mBiasWeights.count;
int nbInputChannels = mKernelWeights.count /
nbOutputChannels;
float kONE = 1.0f, kZERO = 0.0f;
cublasSetStream(mCublas, stream);
cudnnSetStream(mCudnn, stream);
CHECK(cublasSgemm(mCublas, CUBLAS_OP_T, CUBLAS_OP_N,nbOutputChannels, batchSize, nbInputChannels, &kONE,reinterpret_cast<constfloat*>(mKernelWeights.values),
nbInputChannels,reinterpret_cast<const float*>(inputs[0]),nbInputChannels, &kZERO,reinterpret_cast<float*>(outputs[0]),nbOutputChannels));
CHECK(cudnnSetTensor4dDescriptor(mSrcDescriptor,CUDNN_TENSOR_NCHW, CUDNN_DATA_FLOAT, 1, nbOutputChannels, 1, 1));
CHECK(cudnnSetTensor4dDescriptor(mDstDescriptor,CUDNN_TENSOR_NCHW, CUDNN_DATA_FLOAT, batchSize, nbOutputChannels,1, 1));
CHECK(cudnnAddTensor(mCudnn, &kONE, mSrcDescriptor,mBiasWeights.values, &kONE, mDstDescriptor, outputs[0]));
return 0;
}

3.8.8. Serialization

层参数可以和剩下的网络被串连起来。这个串连系统调用了下面的函数：

virtual size_t getSerializationSize() override
{
// 3 integers (number of input channels, number of output channels, bias size), and then the weights:
return sizeof(int)*3 + mKernelWeights.count*sizeof(float) +mBiasWeights.count*sizeof(float);
}

virtual void serialize(void* buffer) override
{
char* d = reinterpret_cast<char*>(buffer), *a = d;
write(d, mNbInputChannels);
write(d, mNbOutputChannels);
write(d, (int)mBiasWeights.count);
serializeFromDevice(d, mKernelWeights);
serializeFromDevice(d, mBiasWeights);
assert(d == a + getSerializationSize());
}
串连被下面的构造函数实现
Deserialization is implemented with the following constructor:

// create the plugin at runtime from a byte stream
FCPlugin(const void* data, size_t length)
{
const char* d = reinterpret_cast<const char*>(data), *a = d;
mNbInputChannels = read<int>(d);
mNbOutputChannels = read<int>(d);
int biasCount = read<int>(d);
mKernelWeights = deserializeToDevice(d, mNbInputChannels *
mNbOutputChannels);
mBiasWeights = deserializeToDevice(d, biasCount);
assert(d == a + length);
}

3.8.9. 添加plugin到网络中（Adding The Plugin Into A Network）

有三种方法可以添加plugin到网络中：

1在定义网络的时候使用INetwork::addPlugin()

2通过parser创建（create）网络

3在网络创建后串转并（De-serialize）

为了使用addPlugin()，请查看TensorRT API（http://docs.nvidia.com/deeplearning/sdk/tensorrt-api/index.html）

3.8.9.1从NvCaffeParser创建plugins

通过NvCaffeParser添加自定义层，通过实现nvcaffeParser::IPluginFactory接口创建一个factory，然后传一个实例到ICaffeParser::parse()。

createPlugin()方法接收到层的名字，然后权重的集合（set）从NVCaffe模型文件中提取出来。然后会将这个集合传到层的构造函数中。名字可以用来区分不同的复杂的plugins。
现在一般没办法提取参数除了从NVCaffe网络描述中提取。因此，这些参数必须要在factory中被指定。
bool isPlugin(const char* name) override
{
return !strcmp(name, "ip2");
}

virtual nvinfer1::IPlugin* createPlugin(const char* layerName,const nvinfer1::Weights* weights, int nbWeights) override
{
// there's no way to pass parameters through from the model definition, so we have to define it here explicitly

static const int NB_OUTPUT_CHANNELS = 10;
assert(isPlugin(layerName) && nbWeights == 2 &&weights[0].type ==DataType::kFLOAT && weights[1].type == DataType::kFLOAT);
assert(mPlugin.get() == nullptr);
mPlugin = std::unique_ptr<FCPlugin>(new FCPlugin(weights,nbWeights, NB_OUTPUT_CHANNELS));
return mPlugin.get();
}

3.8.9.2在Runtime阶段创建plugins
为了整合自定义层在runtime中，实现nvinfer1::IPlugin接口然后将一个factory实例到IInferRuntime::deserializeCudaEngine()。

IPlugin* createPlugin(const char* layerName, const void*serialData,size_t serialLength) override
{
assert(isPlugin(layerName));
assert(mPlugin.get() == nullptr);
mPlugin = std::make_unique<FCPlugin>(serialData,serialLength);
return mPlugin.get();
}

当构造使用NvCaffeParser or deserialized（串连）在runtime阶段，层的实现可能假设数据通过权重（来自NvCaffeParser）或者二进制流（在runtime阶段）将会存在
直到调用initialize()，允许数据在这个函数中被复制到GPU中。