kint8 量化模型封装代码

sunty

2022-03-21 d0a24896f95b4e060011852f80048ebfb0bf5f55

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
#ifndef CHUNK_H_
#define CHUNK_H_
 
//#include "NvInfer.h"
//#include "NvInferPlugin.h"
//#include "NvInferRuntimeCommon.h"
//#include <cuda_runtime.h>
//#include <iostream>
//#include <memory>
//#include <sstream>
//#include <string>
//#include <cassert>
//#include <vector>
 
#include <string>
#include <vector>
#include "NvInfer.h"
 
namespace nvinfer1
{
    class Chunk : public IPluginV2
    {
    public:
        Chunk();
        Chunk(const void* buffer, size_t length);
        ~Chunk();
        int getNbOutputs()const noexcept override;
        Dims getOutputDimensions(int index, const Dims* inputs, int nbInputDims)noexcept  override;
        int initialize()noexcept  override;
        void terminate()noexcept  override;
        size_t getWorkspaceSize(int maxBatchSize) const noexcept  override;
    //    int enqueue(int batchSize, const void* const* inputs, void** outputs, void* workspace, cudaStream_t stream)noexcept override;
 
        int enqueue(int batchSize, const void* const* inputs, void** outputs, void* workspace,
            cudaStream_t stream) noexcept override;
 
        size_t getSerializationSize() const noexcept  override;
        void serialize(void* buffer) const noexcept  override;
        const char* getPluginType() const noexcept  override;
        const char* getPluginVersion() const noexcept  override;
        void destroy()noexcept  override;
        void setPluginNamespace(const char* pluginNamespace)noexcept  override;
        const char* getPluginNamespace() const noexcept  override;
        DataType getOutputDataType(int index, const nvinfer1::DataType* inputTypes, int nbInputs) const noexcept;
        bool isOutputBroadcastAcrossBatch(int outputIndex, const bool* inputIsBroadcasted, int nbInputs) const noexcept;
        bool canBroadcastInputAcrossBatch(int inputIndex) const noexcept;
        void attachToContext(
            cudnnContext* cudnnContext, cublasContext* cublasContext, IGpuAllocator* gpuAllocator);
        void configurePlugin(const PluginTensorDesc* in, int nbInput, const PluginTensorDesc* out, int nbOutput);
        void detachFromContext();
        bool supportsFormatCombination(int pos, const PluginTensorDesc* inOut, int nbInputs, int nbOutputs) const
        {
            return inOut[pos].format == TensorFormat::kLINEAR && inOut[pos].type == DataType::kFLOAT;
        }
        IPluginV2* clone() const noexcept  override;
        bool supportsFormat(DataType type, PluginFormat format) const noexcept override;
        void configureWithFormat(const Dims* inputDims, int nbInputs, const Dims* outputDims, int nbOutputs, DataType type, PluginFormat format, int maxBatchSize) noexcept override;
 
    private:
        std::string _s_plugin_namespace;
        int _n_size_split;
    };
 
    class ChunkPluginCreator : public IPluginCreator
    {
    public:
        ChunkPluginCreator();
        ~ChunkPluginCreator() override = default;
        const char* getPluginName()const noexcept  override;
        const char* getPluginVersion() const noexcept  override;
        const PluginFieldCollection* getFieldNames()noexcept  override;
        IPluginV2* createPlugin(const char* name, const PluginFieldCollection* fc)noexcept  override;
        IPluginV2* deserializePlugin(const char* name, const void* serialData, size_t serialLength)noexcept  override;
        void setPluginNamespace(const char* libNamespace)noexcept  override;
        const char* getPluginNamespace() const noexcept  override;
    private:
        std::string _s_name_space;
        static PluginFieldCollection _fc;
        static std::vector<PluginField> _vec_plugin_attributes;
    };
 
}//nampespace nvinfer1
 
 
#endif