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

公司需求:跑通模型并输出前馈时间。
本文内容:修改官方提供的trtexec(用于前馈时间基准测试的demo),寻找、修改自定义的upsample层,把该层嵌入到trtexec中,然后输入工程的网络结构,返回前馈时间。
接到的网络,TensorRT无法识别的只有upsample层和最后的detection层,主要需要处理的是upsample自定义层。查询了一些资料,在github上下载了一些工程,跑通了一个检测前馈时间的实例(还没有进行准确度测试)。
upsample层是借鉴于github上的工程TensorRT Wrapper

文章目录

  • 结构简介
  • PluginFactory.h
  • upsample.h
  • upsample.cpp
  • upsample.cu
    • enqueue
    • forwardGpu
    • upscale
      • 核函数的调用
      • upscale核函数的原理
    • translate_idx

结构简介

TensorRT Wrapper是依附于yolov3工程的,内含了yolov3层和upsample层的源码,下面主要介绍upsample层。
不管是在trtexec中使用自定义层,还是在其他demo中使用自定义层,都要调用PluginFactory这个类,利用这个类去实例化Plugin(也就是所说的自定义层,详见TensorRT的Plugin)。
  所以想要在trtexec中加入自定义的upsample,主要需要改动的就是PluginFactory.h, upsample.h, upsample.cpp, upsample.cu四个文件,下面依次介绍。

PluginFactory.h

首先介绍PluginFactory头文件的改动,PluginFactory.h 里面内含PluginFactory类,继承于nvinfer1::IPluginFactory和nvcaffeparser1::IPluginFactoryExt类,负责对Plugin层进行实例化。
主要代码:

class PluginFactory : public nvinfer1::IPluginFactory, public nvcaffeparser1::IPluginFactoryExt
{
public:
    inline bool isUpsample(const char* layerName)
    {
        return std::regex_match(layerName , std::regex(R"(layer(\d*)-upsample)"));
    }


    virtual nvinfer1::IPlugin* createPlugin(const char* layerName, const nvinfer1::Weights* weights, int nbWeights) override
    {
        assert(isPlugin(layerName));
        if (isUpsample(layerName))
        {
            // 对于Upsample层来说,输入的Weight的种类数量为0,所以指针为nullptr
            assert(nbWeights == 0 && weights == nullptr);
            //
            mPluginUpsample.emplace_back(std::unique_ptr<UpsampleLayerPlugin>(new UpsampleLayerPlugin(UPSAMPLE_SCALE,CUDA_THREAD_NUM)));
            return mPluginUpsample.back().get();
        }

        else
        {
            assert(0);
            return nullptr;
        }
    }

    nvinfer1::IPlugin* createPlugin(const char* layerName, const void* serialData, size_t serialLength) override
    {
        assert(isPlugin(layerName));

        if (isUpsample(layerName))
        {
            mPluginUpsample.emplace_back(std::unique_ptr<UpsampleLayerPlugin>(new UpsampleLayerPlugin(serialData, serialLength)));
            return mPluginUpsample.back().get();
        }
        else
        {
            assert(0);
            return nullptr;
        }
    }


    bool isPlugin(const char* name) override
    {
        return isPluginExt(name);
    }

    bool isPluginExt(const char* name) override
    {
        //std::cout << "check plugin " << name  << isYolo(name)<< std::endl;
        return isUpsample(name);
    }

    // The application has to destroy the plugin when it knows it's safe to do so.
    void destroyPlugin()
    {
        for (auto& item : mPluginUpsample)
            item.reset();

    }

    void (*nvPluginDeleter)(INvPlugin*){[](INvPlugin* ptr) { if(ptr) ptr->destroy(); }};

    std::vector<std::unique_ptr<UpsampleLayerPlugin>> mPluginUpsample{};
};

这个PluginFactory类和官方demo: samplePlugin中定义的类,结构大致相同,即:isPlugin(), isPluginExt(), 两个createPlugin(),destroyPlugin()。
不同的地方的是

  1. 官方只加一个FC层,所以最后只是声明了一个智能指针。

std::unique_ptr< FCPlugin> mPlugin{nullptr};

但是,在这里,为了能让upsample层的代码可以为多个层服务,所以声明了unique_ptr的一个vector。

std::vector> mPluginUpsample{};

  1. 加入了一个destroy过程,具体我自己还没太弄清楚,挖坑。

std::vector> mPluginUpsample{};

PluginFactory的改动重点就是,智能指针的vector的设置配合isUpsample方法的正则匹配进行多个upsample层的匹配,如下代码所示;又因为这里设置的scale都为2,且为全局变量,所以没有weights需要传入,所以nbWeights == 0,weights==nullptr,其余部分和samplePlugin这个demo中写的几乎一样。

//isUpsample方法中利用正则匹配
return std::regex_match(layerName , std::regex(R"(layer(\d*)-upsample)"));
//最后建立一个vector来存所有指向各个upsample层的unique_ptr指针。
std::vector> mPluginUpsample{};
// 不管是在build还是runtime的plugin实例化过程中(调用两个createPlugin方法),都直接在vector中直接emplace_back即可。
mPluginUpsample.emplace_back(std::unique_ptr< UpsampleLayerPlugin>(new UpsampleLayerPlugin(UPSAMPLE_SCALE,CUDA_THREAD_NUM)));
mPluginUpsample.emplace_back(std::unique_ptr< UpsampleLayerPlugin>(new UpsampleLayerPlugin(serialData, serialLength)));

upsample.h

upsample层的头文件依照tensorRT官方给的模板,声明需要的函数。

  1. 构造、析构函数
  2. getNbOutputs() 返回输出的数量
  3. getOutputDimensions() 返回输出的维度(Dims数据结构)
  4. supportsFormat()和configureWithFormat() 定义datatype和format
  5. getSerializationSize()和serialize() 获得串行化的字节长度的方法和串行化方法
  6. getWorkspaceSize() 获得workSpace需要的size的方法
  7. initialize,terminate和enqueue方法 层的核心方法
  8. forwardGpu方法的声明,为upsample层进行计算的方法。
  9. private中的声明:
        nvinfer1::Dims mCHW;
        DataType mDataType{DataType::kFLOAT};
        float mScale;
        int mOutputWidth;
        int mOutputHeight;
        int mThreadCount;

.
  这里的结构,除了多了forwardGpu之外,和官方给的samplePlugin是一样的,缺少了deserialize相关的方法,因为这个程序中不涉及到并行化问题。
  其中private中的参数mTreadCount是线程数量,构造函数在构造过程中默认该参数值为512,用于后来调用cuda核函数的时候,计算grid和block的设置(grid和block的解释在下面)。

upsample.cpp

下面是对应函数的介绍
1. 构造函数与析构函数

    UpsampleLayerPlugin::UpsampleLayerPlugin(const float scale, const int cudaThread /*= 512*/)
            : mScale(scale),mThreadCount(cudaThread)
    {
    }

    UpsampleLayerPlugin::~UpsampleLayerPlugin()
    {
    }

    UpsampleLayerPlugin::UpsampleLayerPlugin(const void* data, size_t length)
    {
        using namespace Tn;
        const char *d = reinterpret_cast<const char *>(data), *a = d;
        read(d, mCHW);
        read(d, mDataType);
        read(d, mScale);
        read(d, mOutputWidth);
        read(d, mOutputHeight);
        read(d, mThreadCount);
        //std::cout << "read:" << a << " " << mOutputWidth<< " " <
        assert(d == a + length);
    }

build期间的构造函数只是给scale和cudaThread进行了赋值。
runtime期间的构造函数(为了从byte流中创建一个plugin的构造函数),通过read函数读取byte流中对应位置存储的参数值。
析构函数为空。

2. getNbOutputs()
return 1;
3. getOutputDimensions()

   Dims UpsampleLayerPlugin::getOutputDimensions(int index, const Dims* inputs, int nbInputDims)
    {
        mCHW = inputs[0];
        mOutputHeight = inputs[0].d[1]* mScale;
        mOutputWidth = inputs[0].d[2]* mScale;
        return Dims3(mCHW.d[0], mOutputHeight, mOutputWidth);
    }

inputs[0]: 第0个Dims类型的输入(因为upsample层的输入也只有一个层,所以只有input[0])
inputs[0].d[0]/d[1]/d[2]: 分别是C,H,W
return 的是dims3类型,保存的是1个输出层的信息。

4. supportsFormat()和configureWithFormat()

  bool supportsFormat(DataType type, PluginFormat format) const override 
   {
           return (type == DataType::kFLOAT || type == DataType::kHALF || type == DataType::kINT8 )&& format == PluginFormat::kNCHW;
   }
  void UpsampleLayerPlugin::configureWithFormat(const Dims* inputDims, int nbInputs, const Dims* outputDims, int nbOutputs, DataType type, PluginFormat format, int maxBatchSize)
   {
       assert((type == DataType::kFLOAT || type == DataType::kHALF || type == DataType::kINT8) && format == PluginFormat::kNCHW);
       mDataType = type;
   }

这个工程么有设置fp16和int8的选项,只支持,DataType:kFLOAT 和 PluginFormat::KNCHW

5. getSerializationSize()和serialize()

        virtual size_t getSerializationSize() override
       {
           return sizeof(nvinfer1::Dims) + sizeof(mDataType) + sizeof(mScale)
                  + sizeof(mOutputWidth) + sizeof(mOutputHeight) + sizeof(mThreadCount);
       }

计算需要的所有参数需要的总长度,其中包括一个Dims(getNbOutputDimensions返回的输出层的Dims)+ mDataType(kFLOAT)+ mScale +mOutputWidth+mOutputHeight+mThreadCount,即所有执行层运算的参数的总的byte流长度。

  void UpsampleLayerPlugin::serialize(void* buffer)
  {
      using namespace Tn;
      char* d = static_cast<char*>(buffer), *a = d;
      write(d, mCHW);
      write(d, mDataType);
      write(d, mScale);
      write(d, mOutputWidth);
      write(d, mOutputHeight);
      write(d, mThreadCount);
      assert(d == a + getSerializationSize());
  }

按照长度把参数信息写进byte流中。最后的assert是确保写入长度正确。

6. getWorkspaceSize()
这个坑看明白后再填
7. initialize,terminate和enqueue方法 层的核心方法

   int UpsampleLayerPlugin::initialize()
   {
       int inputHeight = mCHW.d[1];
       int inputWidth = mCHW.d[2];

       mOutputHeight = inputHeight * mScale;
       mOutputWidth = inputWidth * mScale;

       return 0;
   }

initialize() 为UpsampleLayerPlugin类的几个私有成员变量赋值。
terminate()为空
enqueue()方法写在了upsample.cu文件中

8. forwardGpu()方法也写在了upsample.cu的文件中

upsample.cu

.cu文件中主要是涉及到了cuda核函数的调用,所以enqueue()和forwardGpu方法都写在这个文件中。

enqueue

enqueue是TensorRT自定义层的核心方法,用来调用cuda核函数或者cuda handle对自定义层进行运算。

int UpsampleLayerPlugin::enqueue(int batchSize, const void* const* inputs, void** outputs, void* workspace, cudaStream_t stream)
    {
        const int channels = mCHW.d[0];
        const int64_t in_height = mCHW.d[1];
        const int64_t in_width = mCHW.d[2];
        const int64_t out_height = mOutputHeight;
        const int64_t out_width = mOutputWidth;
        int totalElems = batchSize * in_height * in_width * channels;

        if (out_height == in_height && out_width == in_width) {
            CUDA_CHECK(cudaMemcpyAsync(outputs[0], inputs[0], totalElems * type2size(mDataType), cudaMemcpyDeviceToDevice, stream));
            CUDA_CHECK(cudaStreamSynchronize(stream));
            return 0;
        }
        switch (mDataType)
        {
            case DataType::kFLOAT :
                forwardGpu<float>((const float *)inputs[0],(float *)outputs[0],batchSize,mCHW.d[0],mOutputHeight,mOutputWidth);
                break;
            case DataType::kHALF:
                forwardGpu<__half>((const __half *)inputs[0],(__half *)outputs[0],batchSize,mCHW.d[0],mOutputHeight,mOutputWidth);
                break;
            case DataType::kINT8:
                forwardGpu<u_int8_t>((const u_int8_t *)inputs[0],(u_int8_t *)outputs[0],batchSize,mCHW.d[0],mOutputHeight,mOutputWidth);
                break;
            default:
                std::cerr << "error data type" << std::endl;
        }
        return 0;
    };

两个CUDA_CHECK还不是很明白,先挖坑。
enqueue函数主要是进行赋值任务和调用forwardGpu()

forwardGpu

forwardGpuforwardGpu主要是对upscale内核函数进行调用

    template <typename Dtype>
    void UpsampleLayerPlugin::forwardGpu(const Dtype* input,Dtype * output, int N,int C,int H ,int W) {
        int numElem = N*C*H*W;
        upscale<<<(numElem + mThreadCount - 1) / mThreadCount, mThreadCount>>>(input,output, numElem, mScale, C, H, W);
    }

其中,函数内部的第二行是对cuda核函数调用的方法。这里我参考这个博文自定义cuda核函数的调用
调用定义的核函数的方式:

kernel<<<1,1>>>(param1,param2,…)

“<<< >>>”中参数的作用是告诉我们该如何启动核函数(比如如何设置线程)。

upscale

核函数的调用

1. 核函数的声明与调用

upscale就是自定义的核函数,规定使用__global__声明核函数

_global_ void kernel(param list){ }

如上面提到的,调用时,如下所示

kernel<<>>(param list);

参数解释:

Dg: int型或者dim3类型(x,y,z)。 用于定义一个grid中的block是如何组织的。 int型则直接表示为1维组织结构。
Db: int型或者dim3类型(x,y,z)。 用于定义一个block中的thread是如何组织的。 int型则直接表示为1维组织结构。
Ns: size_t类型,可缺省,默认为0。 用于设置每个block除了静态分配的共享内存外,最多能动态分配的共享内存大小,单位为byte。 0表示不需要动态分配。
S: cudaStream_t类型,可缺省,默认为0。 表示该核函数位于哪个流。

看完下面的介绍且前文把mThreadCount设置为512可知,一个block中存512个thread,一个grid中存:

(numElem + mThreadCount - 1) / mThreadCount

个block,在这里Dg和Db都是int型,所以grid和block都是一维的。
2. cuda的线程结构

CUDA的线性结构,有三个重要的概念:grid,block,thread
(1) GPU工作时的最小单位是thread
(2) 多个thread可以组成一个block,但是一个block能够包含的thread是有限的。因为一个block上的所有线程最好同时位于同一个处理器核心上,同时共享同一块内存,于是同一块block上的thread可以快速进行同步的动作,而不用担心数据通信壁垒。
(3) 执行相同程序的多个block,可以组成grid。不同block中的thread无法存取同一块共享的内存,无法直接互通或者进行同步。因此,不同block中的thread能合作的程度是比较低的。不过,利用这种模式,可以让程序不用担心显示芯片实际上可以同时执行的thread数目限制。例如,一个具有很少量显示单元的显示芯片,可能会把各个block中的thread顺序执行,而非同时执行。不同的grid可以执行不同的程序。

下面是个一个结构关系图:
TensorRT5.1实践 自定义的Upsample层的方法_第1张图片
另外,block,thread的组织结构可以是二维,三维的,例如上图,block是二维,thread是三维。
CUDA中每一个线程都有一个唯一的标识ThreadIdx,这个ID随着组织结构形式的变化而变化。 (注意:ID的计算,同计算行优先排列的矩阵元素ID思路一样。)


回顾之前我们的矢量加法:
Block是一维的,Tread是一维的:

__global__ void addKernel(int *c, const int *a, const int *b)

    int i = blockIdx.x *blockDim.x + threadIdx.x;  
    c[i] = a[i] + b[i];
}

这里blockIdx.x指的是第blockIdx.x个block,blockDim.x指的是一个block一行具有的thread数(当block为一维的时候,就是一个block具备的thread的数量);threadIdx.x是该block上thread的位置

Block是一维的,thread是二维的:

__global__ void addKernel(int *c, int *a, int *b)
{
    int i = blockIdx.x * blockDim.x * blockDim.y + threadIdx.y * blockDim.x + threadIdx.x;
    c[i] = a[i] + b[i];
}

这里blockDim.x指的是一个block一行具有的thread数量,blockDim.y指的是一个block一列具有的thread数量,threadIdx.y指的是当前thread所在的block的行数,threadIdx.x指的是当前行该thread处于的列数。
Block是二维的,thread是三维的

__global__ void addKernel(int *c, int *a, int *b)
{
    int blockId = blockIdx.x + blockIdx.y * gridDim.x;  
    int i = blockId * (blockDim.x * blockDim.y * blockDim.z)  
        + (threadIdx.z * (blockDim.x * blockDim.y))  
        + (threadIdx.y * blockDim.x) + threadIdx.x; 
    c[i] = a[i] + b[i];
}

blockIdx.y是指这个grid上第blockIdx.y行,gridDim.x指这个grid一行具有的block数量(也即是grid的block列数)
同理。

3. 内存结构
如下图所示,每个 thread 都有自己的一份 register 和 local memory 的空间。同一个 block 中的每个 thread 则有共享的一份 share memory。此外,所有的 thread(包括不同 block 的 thread)都共享一份 global memory、constant memory、和 texture memory。不同的 grid 则有各自的 global memory、constant memory 和 texture memory。

这种特殊的内存结构直接影响着我们的线程分配策略,因为需要通盘考虑资源限制及利用率。
TensorRT5.1实践 自定义的Upsample层的方法_第2张图片
4. 异构编程
如下图所示,是常见的GPU程序的处理流程,其实是一种异构程序,即CPU和GPU的协同。

主机上执行串行代码,设备上则执行并行代码。
TensorRT5.1实践 自定义的Upsample层的方法_第3张图片
以上简单的介绍了cuda计算时thread结构。

upscale核函数的原理

    __global__ void upscale(const Dtype *input, Dtype *output,
                            int no_elements, int scale_factor, int d1, int d2, int d3) {
        int ii = threadIdx.x + blockDim.x * blockIdx.x;
        if (ii >= no_elements) return;  
        int ipidx = translate_idx(ii, d1, d2, d3, scale_factor);
        output[ii]=input[ipidx];
    }

自定义的upscale核函数的代码如上所示,上面由upscale的定义

upscale<<<(numElem + mThreadCount - 1) / mThreadCount, mThreadCount>>>(input,output, numElem, mScale, C, H, W);

可以知道,grid和block都是一维的。
根据上面的示例,我们知道,这个方法先让ii等于现在要使用的这个thread,然后放到translate_idx方法中计算,最后赋值给output。

translate_idx

    __device__ int translate_idx(int ii, int d1, int d2, int d3, int scale_factor) {
        int x, y, z, w;
        w = ii % d3;
        ii = ii/d3;
        z = ii % d2;
        ii = ii/d2;
        y = ii % d1;
        ii = ii/d1;
        x = ii;
        w = w/scale_factor;
        z = z/scale_factor;
        d2 /= scale_factor;
        d3 /= scale_factor;
        return (((x*d1+y)*d2)+z)*d3+w;
    }

由upscale方法可知,传入的参数意义如下:
ii:  当前使用的thread(图像中的每个像素的计算都分配一个thread?)
d1: channel数
d2: 输出的高度
d3: 输出的宽度
scale_factor: scale值

以上便是Upsample的实践过程,都是拼拼凑凑起来的,还需要改进。

你可能感兴趣的:(DL,TensorRT)