CenterNet TensorRT works. TensorRT serialization not yet implemented

Signed-oof-by: Davide Sapienza <sapienza.dav@gmail.com>

include/tkDNN/Layer.h

@@ -32,7 +32,7 @@ enum layerType_t {
 class Layer {
 
 public:
-    Layer(Network *net);
+    Layer(Network *net, bool final = false);
     virtual ~Layer();
     virtual layerType_t getLayerType() = 0;
 
@@ -43,6 +43,7 @@ public:
 
     dataDim_t input_dim, output_dim;
     dnnType *dstData;  //where results will be putted
+    bool final;        //if the layer is the final one
 
     std::string getLayerName() {
         layerType_t type = getLayerType();
@@ -80,7 +81,7 @@ class LayerWgs : public Layer {
 
 public:
     LayerWgs(Network *net, int inputs, int outputs, int kh, int kw, int kt,
-             std::string fname_weights, bool batchnorm = false, bool additional_bias = false); 
+             std::string fname_weights, bool batchnorm = false, bool additional_bias = false, bool final = false); 
     virtual ~LayerWgs();
 
     int inputs, outputs;
@@ -160,7 +161,7 @@ class Conv2d : public LayerWgs {
 public:
     Conv2d( Network *net, int out_ch, int kernelH, int kernelW, 
                 int strideH, int strideW, int paddingH, int paddingW,
-                std::string fname_weights, bool batchnorm = false, bool deConv = false);
+                std::string fname_weights, bool batchnorm = false, bool deConv = false, bool final = false);
     virtual ~Conv2d();
     virtual layerType_t getLayerType() { return LAYER_CONV2D; };
 
@@ -217,15 +218,15 @@ public:
     int out_ch;
     int deformableGroup;
     int kernelH, kernelW, strideH, strideW, paddingH, paddingW;
-protected:
-
     dnnType *ones_d1;
     dnnType *ones_d2;
-    cudnnTensorDescriptor_t biasTensorDesc;
     int chunk_dim;
     dnnType *offset, *mask;
     dnnType *output_conv;
 
+protected:
+
+    cudnnTensorDescriptor_t biasTensorDesc;
     void initCUDNN();
 
 };

include/tkDNN/NetworkRT.h

@@ -31,6 +31,7 @@ using namespace nvinfer1;
 #include "pluginsRT/YoloRT.h"
 #include "pluginsRT/UpsampleRT.h"
 //#include "pluginsRT/Int8Calibrator.h"
+#include "pluginsRT/DeformableConvRT.h"
 
 class PluginFactory : IPluginFactory
 {
@@ -85,6 +86,7 @@ public:
     nvinfer1::ILayer* convert_layer(nvinfer1::ITensor *input, Shortcut *l);
     nvinfer1::ILayer* convert_layer(nvinfer1::ITensor *input, Yolo *l);
     nvinfer1::ILayer* convert_layer(nvinfer1::ITensor *input, Upsample *l);
+    nvinfer1::ILayer* convert_layer(nvinfer1::ITensor *input, DeformConv2d *l);
 
     bool serialize(const char *filename);
     bool deserialize(const char *filename);

include/tkDNN/pluginsRT/DeformableConvRT.h

@@ -0,0 +1,80 @@
+#include<cassert>
+#include "../kernels.h"
+
+
+class DeformableConvRT : public IPlugin {
+
+
+
+public:
+	DeformableConvRT(tk::dnn::DeformConv2d *deformable) {
+		this->defRT = deformable;
+	}
+
+	~DeformableConvRT(){
+
+	}
+
+	int getNbOutputs() const override {
+		return 1;
+	}
+
+	Dims getOutputDimensions(int index, const Dims* inputs, int nbInputDims) override {
+		return DimsCHW{defRT->output_dim.c, defRT->output_dim.h, defRT->output_dim.w};
+	}
+
+	void configure(const Dims* inputDims, int nbInputs, const Dims* outputDims, int nbOutputs, int maxBatchSize) override {
+	}
+
+	int initialize() override {
+
+		return 0;
+	}
+
+	virtual void terminate() override {
+	}
+
+	virtual size_t getWorkspaceSize(int maxBatchSize) const override {
+		return 0;
+	}
+
+	virtual int enqueue(int batchSize, const void*const * inputs, void** outputs, void* workspace, cudaStream_t stream) override {
+
+		dnnType *srcData = (dnnType*)reinterpret_cast<const dnnType*>(inputs[0]);
+		dnnType *output_conv = (dnnType*)reinterpret_cast<const dnnType*>(inputs[1]);
+
+		// split conv2d outputs into offset to mask
+		checkCuda(cudaMemcpy(defRT->offset, defRT->output_conv, 2*defRT->chunk_dim*sizeof(dnnType), cudaMemcpyDeviceToDevice)); 
+		checkCuda(cudaMemcpy(defRT->mask, defRT->output_conv + 2*defRT->chunk_dim, defRT->chunk_dim*sizeof(dnnType), cudaMemcpyDeviceToDevice)); 
+		// kernel sigmoide
+		activationSIGMOIDForward(defRT->mask, defRT->mask, defRT->chunk_dim);
+	
+		// deformable convolution
+		dcn_v2_cuda_forward(srcData, defRT->data_d,
+							defRT->bias2_d, defRT->ones_d1,
+							defRT->offset, defRT->mask,
+							reinterpret_cast<dnnType*>(outputs[0]), defRT->ones_d2,
+							defRT->kernelH, defRT->kernelW,
+							defRT->strideH, defRT->strideW,
+							defRT->paddingH, defRT->paddingW,
+							1, 1,
+							defRT->deformableGroup, 
+							defRT->preconv->input_dim.n, defRT->preconv->input_dim.c, defRT->preconv->input_dim.h, defRT->preconv->input_dim.w,
+							defRT->output_dim.n, defRT->output_dim.c, defRT->output_dim.h, defRT->output_dim.w,
+							defRT->chunk_dim);
+
+		return 0;
+	}
+
+
+	virtual size_t getSerializationSize() override {
+		return 0;
+	}
+
+	virtual void serialize(void* buffer) override {
+		char *buf = reinterpret_cast<char*>(buffer);
+	}
+
+	int size;
+	tk::dnn::DeformConv2d *defRT;
+};