TensorRT5.1实践自定义的FC层的方法

最新推荐文章于 2025-05-11 13:57:33 发布

Nonentity_neu

最新推荐文章于 2025-05-11 13:57:33 发布

阅读量1.8k

点赞数

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

本文链接：https://blog.youkuaiyun.com/qq_33120609/article/details/94593643

版权

TensorRT 专栏收录该内容

19 篇文章

订阅专栏

因为FC层的具体实现，在官方的实例samplePlugin中已经给出了，所以主要看一下核心enqueue函数的实现。(学习阶段)

enqueue

FC的enqueue在头文件中是这么定义的

virtual int enqueue(int batchSize, const void* const* inputs, void** outputs, void* workspace, cudaStream_t stream) override;

batchSize: int类型的batchSize
inputs/outputs: 指向inputs和outputs以及workspace的指针(具体怎么存进去的,看到后回来填坑,可能是没开源把)
stream: cudaStream_t是CUstream_st*类型指针,这里也看不到源码,感觉可能是CUDA stream,管理执行单元的并发.

enqueue方法的具体代码:

定义了对应不同精度的float和__half变量
Cublas

cublasSetStream(mCublas, stream);

cublas是CUDA的线性代数库,这里从cublas_v2.h加载, 点进去是从cublas_api.h中加载, cublas_api.h声明了关于线性代数的函数,这一步通过cublasHandle_t类型的mCublas,和参数stream,调用cublasSetStream方法来获取cuBLAS library的流.

(其中,cublasHandle_t类型是cublasContext的typedef,cublasContext类型,是用来产生CUBLAS library的context用的,context可以看做GPU的线程)

所以其实这里,可以理解为将头文件声明的context和传入enqueue函数的创建好了的stream输入,获取cuBlAS library的流

Cudnn同理Cublas,只不过cudnn应该是负责矩阵偏移的,而cublas是负责矩阵点乘.
operation过程:
直接看里面(fp32):

{
       CHECK(cublasSgemm(mCublas, CUBLAS_OP_T, CUBLAS_OP_N, mNbOutputChannels, batchSize, mNbInputChannels, &oneh,
       reinterpret_cast<const float*>(mDeviceKernel), mNbInputChannels,
       reinterpret_cast<const float*>(inputs[0]), mNbInputChannels, &zerof,
       reinterpret_cast<float*>(outputs[0]), mNbOutputChannels));
       }

(1) CHECK: 检测状态
(2) cublasSgemm: 传送门
　cublasSgemm是CUDA的cublas库的矩阵相乘函数.
　cublas的cublasSgemm函数完成C=αop(A)op(B)+βC的计算，当需要计算C=AB时，显然可以直接设置α=1,β=0。其中op操作对决定矩阵是否转置，即决定该矩阵是按照行优先还是列优先。当我们选择CUBLAS_OP_N时表示不转置，按列优先存储；当我们选择CUBLAS_OP_T时表示需要转置，按行优先存储。
(3) mCublas: 为cublas创建的context
(4) CUBLAS_OP_T:　矩阵Ａ，进行转置，按行优先
(5) CUBLAS_OP_N：矩阵Ｂ，不进行转置，按列优先
最后这个函数得到C=AB的结果的转置.(CHECK检查cuda计算是否顺利?).按照上面博文所讲,结果矩阵存在reinterpret_cast< float *>(output[0])
结合下面的矩阵相加操作,结果就得出来了.