优化仓库

孙天宇

2022-07-12 ce9d187fd294cca192a27f52719094e9df7b1b62

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
#ifndef __PLUGIN_LAYER_H__
#define __PLUGIN_LAYER_H__
 
#include <cassert>
#include <cstring>
#include <cudnn.h>
#include <iostream>
#include <memory>
 
#include "NvInferPlugin.h"
 
#define NV_CUDA_CHECK(status)                                                                      \
    {                                                                                              \
        if (status != 0)                                                                           \
        {                                                                                          \
            std::cout << "Cuda failure: " << cudaGetErrorString(status) << " in file " << __FILE__ \
                      << " at line " << __LINE__ << std::endl;                                     \
            abort();                                                                               \
        }                                                                                          \
    }
 
cudaError_t cudaSLayerV3(const void* input, void* output, const uint32_t& batchSize,
    const uint32_t& n_grid_h_, const uint32_t& n_grid_w_,
    const uint32_t& numOutputClasses, const uint32_t& numBBoxes,
    uint64_t outputSize, cudaStream_t stream);
 
class PluginFactory : public nvinfer1::IPluginFactory
{
 
public:
    PluginFactory();
    nvinfer1::IPlugin* createPlugin(const char* layerName, const void* serialData,
                                    size_t serialLength) override;
    bool isPlugin(const char* name);
    void destroy();
 
private:
    static const int m_MaxLeakyLayers = 72;
    static const int m_ReorgStride = 2;
    static constexpr float m_LeakyNegSlope = 0.1f;
    static const int m_NumBoxes = 5;
    static const int m_NumCoords = 4;
    static const int m_NumClasses = 80;
    static const int m_MaxSLayers = 3;
    int m_LeakyReLUCount = 0;
    int m_SLayerCount = 0;
    nvinfer1::plugin::RegionParameters m_RegionParameters{m_NumBoxes, m_NumCoords, m_NumClasses,
                                                          nullptr};
 
    struct INvPluginDeleter
    {
        void operator()(nvinfer1::plugin::INvPlugin* ptr)
        {
            if (ptr)
            {
                ptr->destroy();
            }
        }
    };
    struct IPluginDeleter
    {
        void operator()(nvinfer1::IPlugin* ptr)
        {
            if (ptr)
            {
                ptr->terminate();
            }
        }
    };
    typedef std::unique_ptr<nvinfer1::plugin::INvPlugin, INvPluginDeleter> unique_ptr_INvPlugin;
    typedef std::unique_ptr<nvinfer1::IPlugin, IPluginDeleter> unique_ptr_IPlugin;
 
    unique_ptr_INvPlugin m_ReorgLayer;
    unique_ptr_INvPlugin m_RegionLayer;
    unique_ptr_INvPlugin m_LeakyReLULayers[m_MaxLeakyLayers];
    unique_ptr_IPlugin m_SLayers[m_MaxSLayers];
};
 
class SLayerV3 : public nvinfer1::IPlugin
{
public:
    SLayerV3(const void* data, size_t length);
    SLayerV3(const uint32_t& numBoxes, const uint32_t& numClasses, const uint32_t& grid_h_,const uint32_t &grid_w_);
    int getNbOutputs() const override;
    nvinfer1::Dims getOutputDimensions(int index, const nvinfer1::Dims* inputs,
                                       int nbInputDims) override;
    void configure(const nvinfer1::Dims* inputDims, int nbInputs, const nvinfer1::Dims* outputDims,
                   int nbOutputs, int maxBatchSize) override;
    int initialize() override;
    void terminate() override;
    size_t getWorkspaceSize(int maxBatchSize) const override;
    int enqueue(int batchSize, const void* const* intputs, void** outputs, void* workspace,
                cudaStream_t stream) override;
    size_t getSerializationSize() override;
    void serialize(void* buffer) override;
 
private:
    template <typename T>
    void write(char*& buffer, const T& val)
    {
        *reinterpret_cast<T*>(buffer) = val;
        buffer += sizeof(T);
    }
 
    template <typename T>
    void read(const char*& buffer, T& val)
    {
        val = *reinterpret_cast<const T*>(buffer);
        buffer += sizeof(T);
    }
    uint32_t m_NumBoxes;
    uint32_t m_NumClasses;
    uint32_t m_GridSize;
    uint64_t m_OutputSize;
    uint32_t _n_grid_h;
    uint32_t _n_grid_w;
};
 
#endif // __PLUGIN_LAYER_H__