mmr/tkDNN
Archived
1
0
Fork 0
Code Issues 46 Pull Requests 12 Releases 4 Activity

TRT8 works with almost every nerual network now!!!!(including demo3d)

Browse Source
This commit is contained in:
perseusdg
2021-10-28 23:35:37 +05:30
parent 8c36dd0431
commit c5e66c6bf6
37 changed files with 1042 additions and 700 deletions
Download Patch File Download Diff File Expand all files Collapse all files
include/tkDNN/pluginsRT/DeformableConvRT.h
+35 -16
View File
@@ -8,14 +8,15 @@
#include <tkdnn.h>
namespace nvinfer1 {
class DeformableConvRT : public IPluginV2 {
class DeformableConvRT : public IPluginV2Ext {
public:
DeformableConvRT(int chunk_dim, int kh, int kw, int sh, int sw, int ph, int pw,
int deformableGroup, int i_n, int i_c, int i_h, int i_w,
int o_n, int o_c, int o_h, int o_w,
tk::dnn::DeformConv2d *deformable = nullptr);
int o_n, int o_c, int o_h, int o_w,std::vector<dnnType> data_H,std::vector<dnnType> bias2_H,
std::vector<dnnType> ones_d1_h,std::vector<dnnType> ones_d2_h,std::vector<dnnType> offsetH,std::vector<dnnType> maskH,int height_ones,
int width_ones,int dim_ones);
~DeformableConvRT();
@@ -25,9 +26,6 @@ namespace nvinfer1 {
Dims getOutputDimensions(int index, const Dims *inputs, int nbInputDims) NOEXCEPT override ;
void configureWithFormat(const Dims *inputDims, int nbInputs, const Dims *outputDims, int nbOutputs, DataType type,
PluginFormat format, int maxBatchSize) NOEXCEPT override ;
int initialize() NOEXCEPT override ;
void terminate() NOEXCEPT override ;
@@ -56,11 +54,26 @@ namespace nvinfer1 {
const char *getPluginVersion() const NOEXCEPT override ;
IPluginV2 *clone() const NOEXCEPT override ;
IPluginV2Ext *clone() const NOEXCEPT override ;
DataType getOutputDataType(int index, const nvinfer1::DataType* inputTypes, int nbInputs) const NOEXCEPT override;
void attachToContext(cudnnContext* cudnnContext, cublasContext* cublasContext, IGpuAllocator* gpuAllocator) NOEXCEPT override;
bool isOutputBroadcastAcrossBatch(int outputIndex, const bool* inputIsBroadcasted, int nbInputs) const NOEXCEPT override;
bool canBroadcastInputAcrossBatch(int inputIndex) const NOEXCEPT override;
void configurePlugin (Dims const *inputDims, int32_t nbInputs, Dims const *outputDims,
int32_t nbOutputs, DataType const *inputTypes, DataType const *outputTypes,
bool const *inputIsBroadcast, bool const *outputIsBroadcast, PluginFormat floatFormat,
int32_t maxBatchSize) NOEXCEPT override;
void detachFromContext() NOEXCEPT override;
cublasStatus_t stat;
cublasHandle_t handle;
cublasHandle_t handle{nullptr};
int i_n, i_c, i_h, i_w;
int o_n, o_c, o_h, o_w;
int size;
@@ -73,12 +86,18 @@ namespace nvinfer1 {
int width_ones;
int dim_ones;
dnnType *data_d;
dnnType *bias2_d;
dnnType *ones_d1;
dnnType *offset;
dnnType *mask;
dnnType *ones_d2;
std::vector<dnnType> data_d_v;
std::vector<dnnType> bias2_d_v;
std::vector<dnnType> ones_d1_v;
std::vector<dnnType> offset_v;
std::vector<dnnType> mask_v;
std::vector<dnnType> ones_d2_v;
dnnType* data_d;
dnnType* bias2_d;
dnnType* ones_d1;
dnnType* offset;
dnnType* mask;
dnnType* ones_d2;
// dnnType *input_n;
// dnnType *offset_n;
// dnnType *mask_n;
@@ -99,9 +118,9 @@ namespace nvinfer1 {
const char *getPluginNamespace() const NOEXCEPT override ;
IPluginV2 *deserializePlugin(const char *name, const void *serialData, size_t serialLength) NOEXCEPT override ;
IPluginV2Ext *deserializePlugin(const char *name, const void *serialData, size_t serialLength) NOEXCEPT override ;
IPluginV2 *createPlugin(const char *name, const PluginFieldCollection *fc) NOEXCEPT override ;
IPluginV2Ext *createPlugin(const char *name, const PluginFieldCollection *fc) NOEXCEPT override ;
const char *getPluginName() const NOEXCEPT override ;
include/tkDNN/pluginsRT/FlattenConcatRT.h
+12 -10
View File
@@ -3,7 +3,7 @@
#include <vector>
#include <utils.h>
namespace nvinfer1 {
class FlattenConcatRT : public IPluginV2IOExt {
class FlattenConcatRT : public IPluginV2Ext {
public:
FlattenConcatRT(int c,int h,int w,int rows,int cols) ;
@@ -24,7 +24,7 @@ namespace nvinfer1 {
#if NV_TENSORRT_MAJOR > 7
int enqueue(int batchSize, const void *const *inputs, void *const *outputs, void *workspace, cudaStream_t stream) NOEXCEPT override ;
#elif NV_TENSORRT_MAJOR == 7
#elif NV_TENSORRT_MAJOR <= 7
int32_t enqueue (int32_t batchSize, const void *const *inputs, void **outputs, void *workspace, cudaStream_t stream) override;
#endif
@@ -42,26 +42,28 @@ namespace nvinfer1 {
void setPluginNamespace(const char *pluginNamespace) NOEXCEPT override ;
IPluginV2IOExt *clone() const NOEXCEPT override ;
IPluginV2Ext *clone() const NOEXCEPT override ;
DataType getOutputDataType(int index, const nvinfer1::DataType* inputTypes, int nbInputs) const NOEXCEPT override;
void configurePlugin(const PluginTensorDesc* in, int nbInput, const PluginTensorDesc* out, int nbOutput) NOEXCEPT override;
void attachToContext(cudnnContext* cudnnContext, cublasContext* cublasContext, IGpuAllocator* gpuAllocator) NOEXCEPT override;
bool isOutputBroadcastAcrossBatch(int outputIndex, const bool* inputIsBroadcasted, int nbInputs) const NOEXCEPT override;
bool canBroadcastInputAcrossBatch(int inputIndex) const NOEXCEPT override;
bool supportsFormatCombination(int pos, const PluginTensorDesc* inOut, int nbInputs, int nbOutputs) const NOEXCEPT override;
void configurePlugin (Dims const *inputDims, int32_t nbInputs, Dims const *outputDims,
int32_t nbOutputs, DataType const *inputTypes, DataType const *outputTypes,
bool const *inputIsBroadcast, bool const *outputIsBroadcast, PluginFormat floatFormat,
int32_t maxBatchSize) NOEXCEPT override;
void detachFromContext() NOEXCEPT override;
bool supportsFormat (DataType type, PluginFormat format) const NOEXCEPT override;
int c, h, w;
int rows, cols;
cublasStatus_t stat;
cublasHandle_t handle;
cublasHandle_t handle{nullptr};
private:
std::string mPluginNamespace;
};
@@ -74,9 +76,9 @@ namespace nvinfer1 {
const char *getPluginNamespace() const NOEXCEPT override ;
IPluginV2IOExt *deserializePlugin(const char *name, const void *serialData, size_t serialLength) NOEXCEPT override ;
IPluginV2Ext *deserializePlugin(const char *name, const void *serialData, size_t serialLength) NOEXCEPT override ;
IPluginV2IOExt *createPlugin(const char *name, const PluginFieldCollection *fc) NOEXCEPT override ;
IPluginV2Ext *createPlugin(const char *name, const PluginFieldCollection *fc) NOEXCEPT override ;
const char *getPluginName() const NOEXCEPT override ;
include/tkDNN/pluginsRT/MaxPoolingFixedSizeRT.h
+20 -8
View File
@@ -6,7 +6,7 @@
namespace nvinfer1 {
class MaxPoolFixedSizeRT : public IPluginV2 {
class MaxPoolFixedSizeRT : public IPluginV2Ext {
public:
MaxPoolFixedSizeRT(int c, int h, int w, int n, int strideH, int strideW, int winSize, int padding) ;
@@ -19,9 +19,6 @@ namespace nvinfer1 {
Dims getOutputDimensions(int index, const Dims *inputs, int nbInputDims) NOEXCEPT override ;
void configureWithFormat(const Dims *inputDims, int nbInputs, const Dims *outputDims, int nbOutputs, DataType type,
PluginFormat format, int maxBatchSize) NOEXCEPT override ;
int initialize() NOEXCEPT override ;
void terminate() NOEXCEPT override ;
@@ -31,7 +28,7 @@ namespace nvinfer1 {
#if NV_TENSORRT_MAJOR > 7
int enqueue(int batchSize, const void *const *inputs, void *const *outputs, void *workspace,
cudaStream_t stream) NOEXCEPT override ;
#elif NV_TENSORRT_MAJOR == 7
#elif NV_TENSORRT_MAJOR <= 7
int32_t enqueue (int32_t batchSize, const void *const *inputs, void **outputs, void *workspace, cudaStream_t stream) override;
#endif
@@ -52,7 +49,22 @@ namespace nvinfer1 {
const char *getPluginVersion() const NOEXCEPT override ;
IPluginV2 *clone() const NOEXCEPT override ;
IPluginV2Ext *clone() const NOEXCEPT override ;
DataType getOutputDataType(int index, const nvinfer1::DataType* inputTypes, int nbInputs) const NOEXCEPT override;
void attachToContext(cudnnContext* cudnnContext, cublasContext* cublasContext, IGpuAllocator* gpuAllocator) NOEXCEPT override;
bool isOutputBroadcastAcrossBatch(int outputIndex, const bool* inputIsBroadcasted, int nbInputs) const NOEXCEPT override;
bool canBroadcastInputAcrossBatch(int inputIndex) const NOEXCEPT override;
void configurePlugin (Dims const *inputDims, int32_t nbInputs, Dims const *outputDims,
int32_t nbOutputs, DataType const *inputTypes, DataType const *outputTypes,
bool const *inputIsBroadcast, bool const *outputIsBroadcast, PluginFormat floatFormat,
int32_t maxBatchSize) NOEXCEPT override;
void detachFromContext() NOEXCEPT override;
int n, c, h, w;
@@ -72,9 +84,9 @@ namespace nvinfer1 {
const char *getPluginNamespace() const NOEXCEPT override ;
IPluginV2 *deserializePlugin(const char *name, const void *serialData, size_t serialLength) NOEXCEPT override ;
IPluginV2Ext *deserializePlugin(const char *name, const void *serialData, size_t serialLength) NOEXCEPT override ;
IPluginV2 *createPlugin(const char *name, const PluginFieldCollection *fc) NOEXCEPT override ;
IPluginV2Ext *createPlugin(const char *name, const PluginFieldCollection *fc) NOEXCEPT override ;
const char *getPluginName() const NOEXCEPT override ;
include/tkDNN/pluginsRT/RegionRT.h
+21 -10
View File
@@ -7,10 +7,10 @@
#include <utils.h>
namespace nvinfer1 {
class RegionRT : public IPluginV2 {
class RegionRT : public IPluginV2Ext {
public:
RegionRT(int classes, int coords, int num);
RegionRT(int classes, int coords, int num,int c,int h,int w);
~RegionRT() ;
@@ -20,9 +20,6 @@ namespace nvinfer1 {
Dims getOutputDimensions(int index, const Dims *inputs, int nbInputDims) NOEXCEPT override ;
void configureWithFormat(const Dims *inputDims, int nbInputs, const Dims *outputDims, int nbOutputs, DataType type,
PluginFormat format, int maxBatchSize) NOEXCEPT override ;
int initialize() NOEXCEPT override ;
@@ -32,7 +29,7 @@ namespace nvinfer1 {
#if NV_TENSORRT_MAJOR > 7
int enqueue(int batchSize, const void *const *inputs, void *const *outputs, void *workspace,
cudaStream_t stream) NOEXCEPT override ;
cudaStream_t stream) NOEXCEPT override ;
#elif NV_TENSORRT_MAJOR == 7
int32_t enqueue (int32_t batchSize, const void *const *inputs, void **outputs, void *workspace, cudaStream_t stream) override;
#endif
@@ -54,7 +51,22 @@ namespace nvinfer1 {
bool supportsFormat(DataType type, PluginFormat format) const NOEXCEPT override ;
IPluginV2 *clone() const NOEXCEPT override ;
DataType getOutputDataType(int index, const nvinfer1::DataType* inputTypes, int nbInputs) const NOEXCEPT override;
void attachToContext(cudnnContext* cudnnContext, cublasContext* cublasContext, IGpuAllocator* gpuAllocator) NOEXCEPT override;
bool isOutputBroadcastAcrossBatch(int outputIndex, const bool* inputIsBroadcasted, int nbInputs) const NOEXCEPT override;
bool canBroadcastInputAcrossBatch(int inputIndex) const NOEXCEPT override;
void configurePlugin (Dims const *inputDims, int32_t nbInputs, Dims const *outputDims,
int32_t nbOutputs, DataType const *inputTypes, DataType const *outputTypes,
bool const *inputIsBroadcast, bool const *outputIsBroadcast, PluginFormat floatFormat,
int32_t maxBatchSize) NOEXCEPT override;
void detachFromContext() NOEXCEPT override;
IPluginV2Ext *clone() const NOEXCEPT override ;
int c, h, w;
int classes, coords, num;
@@ -76,9 +88,9 @@ namespace nvinfer1 {
const char *getPluginNamespace() const NOEXCEPT override ;
IPluginV2 *deserializePlugin(const char *name, const void *serialData, size_t serialLength) NOEXCEPT override ;
IPluginV2Ext *deserializePlugin(const char *name, const void *serialData, size_t serialLength) NOEXCEPT override ;
IPluginV2 *createPlugin(const char *name, const PluginFieldCollection *fc) NOEXCEPT override ;
IPluginV2Ext *createPlugin(const char *name, const PluginFieldCollection *fc) NOEXCEPT override ;
const char *getPluginName() const NOEXCEPT override ;
@@ -96,4 +108,3 @@ namespace nvinfer1 {
};
#endif
include/tkDNN/pluginsRT/ReorgRT.h
+20 -9
View File
@@ -4,10 +4,10 @@
#include <vector>
namespace nvinfer1 {
class ReorgRT : public IPluginV2 {
class ReorgRT : public IPluginV2Ext {
public:
ReorgRT(int stride);
ReorgRT(int stride,int c,int h,int w);
~ReorgRT();
@@ -17,10 +17,6 @@ namespace nvinfer1 {
Dims getOutputDimensions(int index, const Dims *inputs, int nbInputDims) NOEXCEPT override;
void
configureWithFormat(const Dims *inputDims, int nbInputs, const Dims *outputDims, int nbOutputs, DataType type,
PluginFormat format, int maxBatchSize) NOEXCEPT override;
int initialize() NOEXCEPT override;
void terminate() NOEXCEPT override;
@@ -51,7 +47,22 @@ namespace nvinfer1 {
void setPluginNamespace(const char *pluginNamespace) NOEXCEPT override;
IPluginV2 *clone() const NOEXCEPT override;
IPluginV2Ext *clone() const NOEXCEPT override;
DataType getOutputDataType(int index, const nvinfer1::DataType* inputTypes, int nbInputs) const NOEXCEPT override;
void attachToContext(cudnnContext* cudnnContext, cublasContext* cublasContext, IGpuAllocator* gpuAllocator) NOEXCEPT override;
bool isOutputBroadcastAcrossBatch(int outputIndex, const bool* inputIsBroadcasted, int nbInputs) const NOEXCEPT override;
bool canBroadcastInputAcrossBatch(int inputIndex) const NOEXCEPT override;
void configurePlugin (Dims const *inputDims, int32_t nbInputs, Dims const *outputDims,
int32_t nbOutputs, DataType const *inputTypes, DataType const *outputTypes,
bool const *inputIsBroadcast, bool const *outputIsBroadcast, PluginFormat floatFormat,
int32_t maxBatchSize) NOEXCEPT override;
void detachFromContext() NOEXCEPT override;
int c, h, w, stride;
private:
@@ -66,9 +77,9 @@ namespace nvinfer1 {
const char *getPluginNamespace() const NOEXCEPT override;
IPluginV2 *deserializePlugin(const char *name, const void *serialData, size_t serialLength) NOEXCEPT override;
IPluginV2Ext *deserializePlugin(const char *name, const void *serialData, size_t serialLength) NOEXCEPT override;
IPluginV2 *createPlugin(const char *name, const PluginFieldCollection *fc) NOEXCEPT override;
IPluginV2Ext *createPlugin(const char *name, const PluginFieldCollection *fc) NOEXCEPT override;
const char *getPluginName() const NOEXCEPT override;
include/tkDNN/pluginsRT/ReshapeRT.h
+21 -8
View File
@@ -9,10 +9,10 @@ using namespace tk::dnn;
namespace nvinfer1 {
class ReshapeRT : public IPluginV2 {
class ReshapeRT : public IPluginV2Ext {
public:
explicit ReshapeRT(dataDim_t newDim) ;
ReshapeRT(int n,int c,int h,int w) ;
ReshapeRT(const void *data, size_t length) ;
@@ -22,8 +22,6 @@ namespace nvinfer1 {
Dims getOutputDimensions(int index, const Dims *inputs, int nbInputDims) NOEXCEPT override ;
void configureWithFormat(const Dims *inputDims, int nbInputs, const Dims *outputDims, int nbOutputs, DataType type,
PluginFormat format, int maxBatchSize) NOEXCEPT override ;
int initialize() NOEXCEPT override ;
@@ -53,9 +51,24 @@ namespace nvinfer1 {
void setPluginNamespace(const char *pluginNamespace) NOEXCEPT override ;
IPluginV2 *clone() const NOEXCEPT override ;
IPluginV2Ext *clone() const NOEXCEPT override ;
DataType getOutputDataType(int index, const nvinfer1::DataType* inputTypes, int nbInputs) const NOEXCEPT override;
void attachToContext(cudnnContext* cudnnContext, cublasContext* cublasContext, IGpuAllocator* gpuAllocator) NOEXCEPT override;
bool isOutputBroadcastAcrossBatch(int outputIndex, const bool* inputIsBroadcasted, int nbInputs) const NOEXCEPT override;
bool canBroadcastInputAcrossBatch(int inputIndex) const NOEXCEPT override;
void configurePlugin (Dims const *inputDims, int32_t nbInputs, Dims const *outputDims,
int32_t nbOutputs, DataType const *inputTypes, DataType const *outputTypes,
bool const *inputIsBroadcast, bool const *outputIsBroadcast, PluginFormat floatFormat,
int32_t maxBatchSize) NOEXCEPT override;
void detachFromContext() NOEXCEPT override;
int n, c, h, w;
dataDim_t new_dim;
private:
std::string mPluginNamespace;
};
@@ -68,9 +81,9 @@ namespace nvinfer1 {
const char *getPluginNamespace() const NOEXCEPT override ;
IPluginV2 *deserializePlugin(const char *name, const void *serialData, size_t serialLength) NOEXCEPT override ;
IPluginV2Ext *deserializePlugin(const char *name, const void *serialData, size_t serialLength) NOEXCEPT override ;
IPluginV2 *createPlugin(const char *name, const PluginFieldCollection *fc) NOEXCEPT override ;
IPluginV2Ext *createPlugin(const char *name, const PluginFieldCollection *fc) NOEXCEPT override ;
const char *getPluginName() const NOEXCEPT override ;
include/tkDNN/pluginsRT/ResizeLayerRT.h
+24 -12
View File
@@ -6,10 +6,10 @@
namespace nvinfer1 {
class ResizeLayerRT : public IPluginV2 {
class ResizeLayerRT : public IPluginV2Ext {
public:
ResizeLayerRT(int c, int h, int w) ;
ResizeLayerRT(int oc, int oh, int ow,int ic,int ih,int iw) ;
ResizeLayerRT(const void *data, size_t length) ;
@@ -19,10 +19,6 @@ namespace nvinfer1 {
Dims getOutputDimensions(int index, const Dims *inputs, int nbInputDims) NOEXCEPT override ;
void configureWithFormat(const Dims *inputDims, int nbInputs, const Dims *outputDims, int nbOutputs, DataType type,
PluginFormat format, int maxBatchSize) NOEXCEPT override ;
int initialize() NOEXCEPT override ;
void terminate() NOEXCEPT override ;
@@ -32,12 +28,10 @@ namespace nvinfer1 {
#if NV_TENSORRT_MAJOR > 7
int enqueue(int batchSize, const void *const *inputs, void *const *outputs, void *workspace,
cudaStream_t stream) NOEXCEPT override ;
#elif NV_TENSORRT_MAJOR == 7
#elif NV_TENSORRT_MAJOR <= 7
int32_t enqueue (int32_t batchSize, const void *const *inputs, void **outputs, void *workspace, cudaStream_t stream) override;
#endif
size_t getSerializationSize() const NOEXCEPT override ;
void serialize(void *buffer) const NOEXCEPT override ;
@@ -54,7 +48,22 @@ namespace nvinfer1 {
void setPluginNamespace(const char *pluginNamespace) NOEXCEPT override ;
IPluginV2 *clone() const NOEXCEPT override ;
IPluginV2Ext *clone() const NOEXCEPT override ;
DataType getOutputDataType(int index, const nvinfer1::DataType* inputTypes, int nbInputs) const NOEXCEPT override;
void attachToContext(cudnnContext* cudnnContext, cublasContext* cublasContext, IGpuAllocator* gpuAllocator) NOEXCEPT override;
bool isOutputBroadcastAcrossBatch(int outputIndex, const bool* inputIsBroadcasted, int nbInputs) const NOEXCEPT override;
bool canBroadcastInputAcrossBatch(int inputIndex) const NOEXCEPT override;
void configurePlugin (Dims const *inputDims, int32_t nbInputs, Dims const *outputDims,
int32_t nbOutputs, DataType const *inputTypes, DataType const *outputTypes,
bool const *inputIsBroadcast, bool const *outputIsBroadcast, PluginFormat floatFormat,
int32_t maxBatchSize) NOEXCEPT override;
void detachFromContext() NOEXCEPT override;
int i_c, i_h, i_w, o_c, o_h, o_w;
@@ -70,9 +79,9 @@ namespace nvinfer1 {
const char *getPluginNamespace() const NOEXCEPT override ;
IPluginV2 *deserializePlugin(const char *name, const void *serialData, size_t serialLength) NOEXCEPT override ;
IPluginV2Ext *deserializePlugin(const char *name, const void *serialData, size_t serialLength) NOEXCEPT override ;
IPluginV2 *createPlugin(const char *name, const PluginFieldCollection *fc) NOEXCEPT override ;
IPluginV2Ext *createPlugin(const char *name, const PluginFieldCollection *fc) NOEXCEPT override ;
const char *getPluginName() const NOEXCEPT override ;
@@ -80,6 +89,9 @@ namespace nvinfer1 {
const PluginFieldCollection *getFieldNames() NOEXCEPT override ;
private:
static PluginFieldCollection mFC;
static std::vector<PluginField> mPluginAttributes;
include/tkDNN/pluginsRT/ShortcutRT.h
+19 -8
View File
@@ -10,10 +10,10 @@
namespace nvinfer1 {
class ShortcutRT : public IPluginV2 {
class ShortcutRT : public IPluginV2Ext {
public:
ShortcutRT(tk::dnn::dataDim_t bdim, bool mul);
ShortcutRT(int bc,int bh,int bw,int c,int h,int w ,bool mul);
~ShortcutRT();
@@ -23,8 +23,19 @@ namespace nvinfer1 {
Dims getOutputDimensions(int index, const Dims *inputs, int nbInputDims) NOEXCEPT override;
void configureWithFormat(const Dims *inputDims, int nbInputs, const Dims *outputDims, int nbOutputs, DataType type,
PluginFormat format, int maxBatchSize) NOEXCEPT override;
void configurePlugin (Dims const *inputDims, int32_t nbInputs, Dims const *outputDims, int32_t nbOutputs,
DataType const *inputTypes, DataType const *outputTypes, bool const *inputIsBroadcast,
bool const *outputIsBroadcast, PluginFormat floatFormat, int32_t maxBatchSize) NOEXCEPT override;
bool isOutputBroadcastAcrossBatch (int32_t outputIndex, bool const *inputIsBroadcasted, int32_t nbInputs) const NOEXCEPT override;
bool canBroadcastInputAcrossBatch (int32_t inputIndex) const NOEXCEPT override;
void attachToContext (cudnnContext *, cublasContext *, IGpuAllocator *) NOEXCEPT override;
void detachFromContext () NOEXCEPT override;
DataType getOutputDataType(int32_t index, nvinfer1::DataType const *inputTypes, int32_t nbInputs) const NOEXCEPT override;
int initialize() NOEXCEPT override;
@@ -56,10 +67,10 @@ namespace nvinfer1 {
void setPluginNamespace(const char *pluginNamespace) NOEXCEPT override;
IPluginV2 *clone() const NOEXCEPT override;
IPluginV2Ext *clone() const NOEXCEPT override;
int c, h, w;
int bc, bh, bw;
int bc, bh, bw,bl;
bool mul;
tk::dnn::dataDim_t bDim;
private:
@@ -75,9 +86,9 @@ namespace nvinfer1 {
const char *getPluginNamespace() const NOEXCEPT override;
IPluginV2 *deserializePlugin(const char *name, const void *serialData, size_t serialLength) NOEXCEPT override;
IPluginV2Ext *deserializePlugin(const char *name, const void *serialData, size_t serialLength) NOEXCEPT override;
IPluginV2 *createPlugin(const char *name, const PluginFieldCollection *fc) NOEXCEPT override;
IPluginV2Ext *createPlugin(const char *name, const PluginFieldCollection *fc) NOEXCEPT override;
const char *getPluginName() const NOEXCEPT override;
include/tkDNN/pluginsRT/UpsampleRT.h
+24 -12
View File
@@ -8,10 +8,10 @@
namespace nvinfer1 {
class UpsampleRT : public IPluginV2 {
class UpsampleRT : public IPluginV2Ext {
public:
explicit UpsampleRT(int stride);
UpsampleRT(int stride,int c,int h,int w);
UpsampleRT(const void *data, size_t length);
@@ -21,9 +21,6 @@ namespace nvinfer1 {
Dims getOutputDimensions(int index, const Dims *inputs, int nbInputDims) NOEXCEPT override;
void configureWithFormat(const Dims *inputDims, int nbInputs, const Dims *outputDims, int nbOutputs, DataType type,
PluginFormat format, int maxBatchSize) NOEXCEPT override;
int initialize() NOEXCEPT override;
void terminate() NOEXCEPT override;
@@ -54,12 +51,27 @@ namespace nvinfer1 {
void setPluginNamespace(const char *pluginNamespace) NOEXCEPT override;
IPluginV2 *clone() const NOEXCEPT override ;
IPluginV2Ext *clone() const NOEXCEPT override ;
int c, h, w, stride;
private:
std::string mPluginNamespace;
};
bool isOutputBroadcastAcrossBatch (int32_t outputIndex, bool const *inputIsBroadcasted, int32_t nbInputs) const NOEXCEPT override;
bool canBroadcastInputAcrossBatch (int32_t inputIndex) const NOEXCEPT override;
void configurePlugin (Dims const *inputDims, int32_t nbInputs, Dims const *outputDims, int32_t nbOutputs,
DataType const *inputTypes, DataType const *outputTypes, bool const *inputIsBroadcast,
bool const *outputIsBroadcast, PluginFormat floatFormat, int32_t maxBatchSize) NOEXCEPT override;
void attachToContext (cudnnContext *, cublasContext *, IGpuAllocator *) NOEXCEPT override;
void detachFromContext () NOEXCEPT override;
DataType getOutputDataType (int32_t index, nvinfer1::DataType const *inputTypes, int32_t nbInputs) const NOEXCEPT override;
int c, h, w, stride;
private:
std::string mPluginNamespace;
};
class UpsampleRTPluginCreator : public IPluginCreator {
public:
@@ -69,9 +81,9 @@ namespace nvinfer1 {
const char *getPluginNamespace() const NOEXCEPT override;
IPluginV2 *deserializePlugin(const char *name, const void *serialData, size_t serialLength) NOEXCEPT override;
IPluginV2Ext *deserializePlugin(const char *name, const void *serialData, size_t serialLength) NOEXCEPT override;
IPluginV2 *createPlugin(const char *name, const PluginFieldCollection *fc) NOEXCEPT override;
IPluginV2Ext *createPlugin(const char *name, const PluginFieldCollection *fc) NOEXCEPT override;
const char *getPluginName() const NOEXCEPT override;
include/tkDNN/pluginsRT/YoloRT.h
+24 -11
View File
@@ -10,10 +10,11 @@
#define YOLORT_CLASSNAME_W 256
namespace nvinfer1 {
class YoloRT : public IPluginV2 {
class YoloRT : public IPluginV2Ext {
public:
YoloRT(int classes, int num, tk::dnn::Yolo *Yolo = nullptr, int n_masks = 3, float scale_xy = 1,
YoloRT(int classes, int num,int c,int h,int w,std::vector<std::string> classNames,
std::vector<float> masks_v,std::vector<float> bias_v, int n_masks = 3, float scale_xy = 1,
float nms_thresh = 0.45, int nms_kind = 0, int new_coords = 0);
YoloRT(const void *data, size_t length);
@@ -25,9 +26,6 @@ namespace nvinfer1 {
Dims getOutputDimensions(int index, const Dims *inputs, int nbInputDims) NOEXCEPT override;
void configureWithFormat(const Dims *inputDims, int nbInputs, const Dims *outputDims, int nbOutputs, DataType type,
PluginFormat format, int maxBatchSize) NOEXCEPT override;
int initialize() NOEXCEPT override;
void terminate() NOEXCEPT override;
@@ -59,9 +57,24 @@ namespace nvinfer1 {
void setPluginNamespace(const char *pluginNamespace) NOEXCEPT override;
IPluginV2 *clone() const NOEXCEPT override;
IPluginV2Ext *clone() const NOEXCEPT override;
DataType getOutputDataType(int index, const nvinfer1::DataType* inputTypes, int nbInputs) const NOEXCEPT override;
void attachToContext(cudnnContext* cudnnContext, cublasContext* cublasContext, IGpuAllocator* gpuAllocator) NOEXCEPT override;
bool isOutputBroadcastAcrossBatch(int outputIndex, const bool* inputIsBroadcasted, int nbInputs) const NOEXCEPT override;
bool canBroadcastInputAcrossBatch(int inputIndex) const NOEXCEPT override;
void configurePlugin (Dims const *inputDims, int32_t nbInputs, Dims const *outputDims,
int32_t nbOutputs, DataType const *inputTypes, DataType const *outputTypes,
bool const *inputIsBroadcast, bool const *outputIsBroadcast, PluginFormat floatFormat,
int32_t maxBatchSize) NOEXCEPT override;
void detachFromContext() NOEXCEPT override;
tk::dnn::Yolo *yolo;
int c, h, w;
int classes, num, n_masks;
float scaleXY;
@@ -71,8 +84,8 @@ namespace nvinfer1 {
int NUM = 0;
std::vector<std::string> classesNames;
dnnType *mask;
dnnType *bias;
std::vector<dnnType> mask;
std::vector<dnnType> bias;
int entry_index(int batch, int location, int entry) {
int n = location / (w * h);
@@ -93,9 +106,9 @@ namespace nvinfer1 {
const char *getPluginNamespace() const NOEXCEPT override;
IPluginV2 *deserializePlugin(const char *name, const void *serialData, size_t serialLength) NOEXCEPT override;
IPluginV2Ext *deserializePlugin(const char *name, const void *serialData, size_t serialLength) NOEXCEPT override;
IPluginV2 *createPlugin(const char *name, const PluginFieldCollection *fc) NOEXCEPT override;
IPluginV2Ext *createPlugin(const char *name, const PluginFieldCollection *fc) NOEXCEPT override;
const char *getPluginName() const NOEXCEPT override;