Added constant and zero padding with ISliceLayer,need to add them for tensorrt versions less than 8.2 using IPluginV2Ext instead of ISliceLayer since they dont seem to support reflect and zero

This commit is contained in:
perseusdg
2022-01-04 09:03:04 +05:30
parent ba022663f1
commit b75cecb105
8 changed files with 396 additions and 4 deletions
+2 -1
View File
@@ -535,12 +535,13 @@ typedef enum {
class Padding : public Layer {
public:
Padding(Network *net,int32_t pad_h,int32_t pad_w,tkdnnPaddingMode_t padding_mode);
Padding(Network *net,int32_t pad_h,int32_t pad_w,tkdnnPaddingMode_t padding_mode,float constant = 0.0);
virtual ~Padding();
virtual layerType_t getLayerType(){return LAYER_PADDING ;};
virtual dnnType* infer(dataDim_t& dim,dnnType* srcData);
int32_t paddingH,paddingW;
tkdnnPaddingMode_t padding_mode;
float constant;
};
+4
View File
@@ -51,4 +51,8 @@ void scalAdd(dnnType* dstData, int size, float alpha, float beta, int inc, cudaS
void reflection_pad2d_out_forward(int32_t pad_h,int32_t pad_w,float *srcData,float *dstData,int32_t input_h,int32_t input_w,int32_t plane_dim,int32_t n_batch,cudaStream_t cudaStream = cudaStream_t(0));
void constant_pad2d_forward(dnnType *srcData,dnnType *dstData,int32_t input_h,int32_t input_w,int32_t output_h,
int32_t output_w,int32_t c,int32_t n,int32_t padT,int32_t padL,dnnType constant,cudaStream_t cudaStream = cudaStream_t(0));
#endif //KERNELS_H
+5 -1
View File
@@ -1,3 +1,6 @@
#ifndef _FLATTENCONCATRT_PLUGIN_H
#define _FLATTENCONCATRT_PLUGIN_H
#include<cassert>
#include <NvInfer.h>
#include <vector>
@@ -93,4 +96,5 @@ namespace nvinfer1 {
};
REGISTER_TENSORRT_PLUGIN(FlattenConcatRTPluginCreator);
};
};
#endif
@@ -0,0 +1,99 @@
#ifndef _REFLECTIONPADDINGRT_PLUGIN_H
#define _REFLECTIONPADDINGRT_PLUGIN_H
#include<cassert>
#include <NvInfer.h>
#include <vector>
#include <utils.h>
#include <kernels.h>
namespace nvinfer1{
class ReflectionPaddingRT : public IPluginV2Ext {
public:
ReflectionPaddingRT(int32_t padH,int32_t padW,int32_t input_h,int32_t input_w,int32_t output_h,int32_t output_w,int32_t c,int32_t n);
ReflectionPaddingRT(const void *data,size_t length);
~ReflectionPaddingRT();
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 ;
#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
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 ;
void destroy() NOEXCEPT override ;
const char *getPluginType() const NOEXCEPT override ;
const char *getPluginVersion() const NOEXCEPT override;
const char *getPluginNamespace() const NOEXCEPT override ;
void setPluginNamespace(const char *pluginNamespace) 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;
bool supportsFormat (DataType type, PluginFormat format) const NOEXCEPT override;
int32_t padH,padW,input_h,input_w,output_h,output_w,n,c;
private:
std::string mPluginNamespace;
};
class ReflectionPaddingRTPluginCreator : public IPluginCreator {
public:
ReflectionPaddingRTPluginCreator();
void setPluginNamespace(const char *pluginNamespace) NOEXCEPT override ;
const char *getPluginNamespace() const NOEXCEPT override ;
IPluginV2Ext *deserializePlugin(const char *name, const void *serialData, size_t serialLength) NOEXCEPT override ;
IPluginV2Ext *createPlugin(const char *name, const PluginFieldCollection *fc) NOEXCEPT override ;
const char *getPluginName() const NOEXCEPT override ;
const char *getPluginVersion() const NOEXCEPT override;
const PluginFieldCollection *getFieldNames() NOEXCEPT override ;
private:
static PluginFieldCollection mFC;
static std::vector<PluginField> mPluginAttributes;
std::string mPluginNamespace;
};
};
#endif