分别使用 TensorRT 和 CUDA 加速 MTCNN

Github 现有的 TensorRT 加速的 MTCNN 【​​PKUZHOU​​/​​MTCNN_FaceDetection_TensorRT​​】不是基于插件的，而是走了使用 scale 和 ReLU 、eltwise-sum 层 “曲线救国”的路线——

PKUZHOU 认为 PReLU 会破坏 TensorRT 的 CBR 优化，但实际上实现 PReLU 插件以后耗时更少，如图

左侧是“曲线救国”版，右侧是实现了 PReLU 插件，一张 1920x1080 的图像，能有200多 ms 的提升（原谅我笔记本的显卡是 GTX 970m，古老的 Maxwell 架构，连半精度都不支持，更别提 Int8 了，抹泪）。

插件的实现代码

prelu_plugn.h

#ifndef PRELU_PLUGIN_H
#define PRELU_PLUGIN_H

#include "kernels.h"

#include <cstring>
#include <assert.h>

#include <cuda.h>
#include <cuda_runtime_api.h>
#include <cuda_fp16.h> // __half

#include <NvInfer.h>
#include <NvCaffeParser.h>

using namespace nvinfer1;
using namespace nvcaffeparser1;

/*
    Prelu layer
*/
class PreluPlugin : public IPlugin
{
public:
    PreluPlugin(const Weights* weights, int nbWeights);
    PreluPlugin(const void* buffer, size_t size);
    ~PreluPlugin();

    Dims getOutputDimensions(int index, const Dims* inputs, int nbInputDims);
    int enqueue(int batchSize, const void* const* inputs, void** outputs, void*, cudaStream_t stream);
    int getNbOutputs() const override
    {
        return 1;
    };

    void configure(const Dims* inputs, int nbInputs, const Dims* outputs, int nbOutputs, int) override;
    void serialize(void* buffer) override;
    size_t getSerializationSize() override;

    inline size_t getWorkspaceSize(int) const override
    {
        return 0;
    }

    int initialize() override;

    void terminate() override;

protected:
    int m_input_c;
    int m_input_h;
    int m_input_w;
    int m_input_count;
    bool m_channel_shared {false};
    Weights m_weights;
    void* m_device_kernel{nullptr};

private:
    void deserializeToDevice(const char*& hostBuffer, void*& deviceWeights, size_t size)
    {
        deviceWeights = copyToDevice(hostBuffer, size);
        hostBuffer += size;
    }

    // 将 host 的 buffer 上的值拷贝到 device (还会开辟设备内存)上
    void* copyToDevice(const void* data, size_t count)
    {
        void* deviceData;
        cudaMalloc(&deviceData, count);
        cudaMemcpy(deviceData, data, count, cudaMemcpyHostToDevice);
        return deviceData;
    }

    template<typename T> void read(const char*& buffer, T& val)
    {
        val = *reinterpret_cast<const T*>(buffer);
        buffer += sizeof(T);
    }

    template<typename T> void write(char*& buffer, const T& val)
    {
        *reinterpret_cast<T*>(buffer) = val;
        buffer += sizeof(T);
    }

    size_t type2size(nvinfer1::DataType type)
    {
//        return sizeof(float);
        return type == nvinfer1::DataType::kFLOAT ? sizeof(float) : sizeof(__half);
    }

    // 将 Weights 的 values 中的值拷贝到 host 的 buffer 上
    void convertAndCopyToBuffer(char*& buffer, const Weights& weights)
    {
        memcpy(buffer, weights.values, weights.count * type2size(weights.type));
        buffer += weights.count * type2size(weights.type);
    }
};



#endif // PRELU_PLUGIN_H

prelu_plugin.cpp

#include "prelu_plugin.h"
#include <iostream>

using namespace nvinfer1;
using namespace nvcaffeparser1;
//using namespace plugin;

PreluPlugin::PreluPlugin(const Weights* weights, int nbWeights)
{
    assert(nbWeights == 1);
    m_weights = weights[0];
    assert(m_weights.type == DataType::kFLOAT || m_weights.type == DataType::kHALF);
    // 为 values 开辟空间
    m_weights.values = malloc(m_weights.count * type2size(m_weights.type));
    // weights[0].values -> m_weights.values
    memcpy(const_cast<void*>(m_weights.values), weights[0].values, m_weights.count * type2size(m_weights.type));
}

PreluPlugin::PreluPlugin(const void* buffer, size_t size)
{
    // 反序列化：和序列化的顺序相同，注意不同的数据类型
    const char* d = reinterpret_cast<const char*>(buffer), *a = d;
    read<int>(d, m_input_c);
    read<int>(d, m_input_h);
    read<int>(d, m_input_w);
    read<int>(d, m_input_count);
    read<bool>(d, m_channel_shared);
    read<int64_t>(d, m_weights.count);
    read<DataType>(d, m_weights.type);

//    m_weights.values = nullptr;
    m_weights.values = malloc(m_weights.count * type2size(m_weights.type));

    //deserializeToDevice(d,m_device_kernel,m_weights.count);
    //  d -> m_weights.values
    memcpy(const_cast<void*>(m_weights.values), d, m_weights.count * type2size(m_weights.type));
    d += m_weights.count * type2size(m_weights.type); // 指针继续向后

    assert(d == a + size);
}

PreluPlugin::~PreluPlugin()
{
//    std::cout << "~PreluPlugin  "<< std::endl;
//    if (m_weights.values)
//    {
//        free(const_cast<void*>(m_weights.values));
//    }
}

// 仅在序列化时调用该方法
Dims PreluPlugin::getOutputDimensions(int index, const Dims* inputs, int nbInputDims)
{
//    std::cout << "0~getOutputDimensions  " << std::endl;
    assert(index == 0 && nbInputDims == 1 && inputs[0].nbDims == 3);
    return DimsCHW(inputs[0].d[0], inputs[0].d[1], inputs[0].d[2]);
}

// 仅在序列化时调用该方法
void PreluPlugin::configure(const Dims* inputs, int nbInputs, const Dims* outputs, int nbOutputs, int)
{
//    std::cout << "1~configure  " << std::endl;
    m_input_c = inputs[0].d[0];
    m_input_h = inputs[0].d[1];
    m_input_w = inputs[0].d[2];
    m_input_count = m_input_c * m_input_h * m_input_w;
}

size_t PreluPlugin::getSerializationSize()
{
    return 4 * sizeof(int) + sizeof(bool) +
            sizeof(m_weights.count)
           + sizeof(m_weights.type)
           +  m_weights.count * type2size(m_weights.type);
}

void PreluPlugin::serialize(void* buffer)
{
    char* d = static_cast<char*>(buffer), *a = d;
    write(d, m_input_c);
    write(d, m_input_h);
    write(d, m_input_w);
    write(d, m_input_count);
    write(d, m_channel_shared);
    write(d, m_weights.count);
    write(d, m_weights.type);

    convertAndCopyToBuffer(d, m_weights);
    assert(d == a + getSerializationSize());
}

int PreluPlugin::enqueue(int batchSize, const void* const* inputs, void** outputs, void*, cudaStream_t stream)
{
    const float* bottom_data = reinterpret_cast<const float*>(inputs[0]);
    float* top_data = reinterpret_cast<float*>(outputs[0]);

    const int count = batchSize * m_input_count;
    const int dim = m_input_h * m_input_w;
    const int channels = m_input_c;
    const int div_factor = m_channel_shared ? channels : 1; //m_channel_shared default is false

    pReLUForward(count, channels, dim, bottom_data, top_data, m_device_kernel, div_factor, stream);

    return 0;
}

int PreluPlugin::initialize()
{
//    std::cout << "2~initialize~0  "<< m_device_kernel << std::endl;
    cudaMalloc(&m_device_kernel, m_weights.count * type2size(m_weights.type));
    cudaMemcpy(m_device_kernel, m_weights.values, m_weights.count * type2size(m_weights.type), cudaMemcpyHostToDevice);
    return 0;
}

// engine 销毁时会调用
void PreluPlugin::terminate()
{
//    std::cout << "~terminate  "<< m_device_kernel << std::endl;

    if (m_weights.values)
    {
        free(const_cast<void*>(m_weights.values));
    }

    if (m_device_kernel)
    {
        cudaFree(m_device_kernel);
        m_device_kernel = nullptr;
    }
}

kernel.cu

__global__ void pReLU(const int n, const int channels, const int dim,
                      const float* in, float* out, const float* slope_data, const int div_factor)
{
    CUDA_KERNEL_LOOP(index, n)
    {
        int c = (index / dim) % channels / div_factor;
        out[index] = in[index] > 0 ? in[index] : in[index] * slope_data[c];
    }
}

void pReLUForward(const int count, const int channels, const int dim, const float* bottom_data,
                         float* top_data, void* mDeviceKernel, const int div_factor, cudaStream_t stream)
{
    pReLU <<< CAFFE_GET_BLOCKS(count), CAFFE_CUDA_NUM_THREADS, 0, stream>>>(count, channels, dim,
                                                                            bottom_data, top_data,
                                                                            static_cast<const float*>(mDeviceKernel),  // slope_data
                                                                            div_factor);

    CUDA_POST_KERNEL_CHECK;
}

继而通过 Valgrind 和 gProf 进行性能分析可以发现，

 

如果优化 nms 和 image2Matrix 方法的话，可以进一步提高性能；

由于多个 Pnet 的检测也是相互独立的，所以还可以使用多线程并行，然后多个流在 GPU （最好支持 HyperQ）上的 Overlap 可以再进一步提高性能。

上图是我没用 TensorRT，直接用原生的 CUDA 加速的效果，迭代 20 次，平均每次仅花费 60 ms 左右。

经过进一步优化纯 CUDA 的算法，一次迭代仅需 44 ms（ 仍然是在 min_size = 30, thresh_p = 0.7, thresh_r = 0.7, thresh_o = 0.7, thresh_nms_p = 0.5, thresh_nms_r = 0.5, thresh_nms_o = 0.5 的条件下）

 如果通过 TensorRT 加速应该会取得更优的性能。

 最后是检测的结果：