CenterNet TensorRT works. TensorRT serialization not yet implemented

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

src/Conv2d.cpp

@@ -119,10 +119,10 @@ void Conv2d::inferCUDNN(dnnType* srcData, bool back) {
 
 Conv2d::Conv2d( Network *net, int out_ch, int kernelH, int kernelW,
                 int strideH, int strideW, int paddingH, int paddingW,
-                std::string fname_weights, bool batchnorm, bool deConv) :
+                std::string fname_weights, bool batchnorm, bool deConv, bool final) :
     
     LayerWgs(net, net->getOutputDim().c, out_ch, kernelH, kernelW, 1, 
-             fname_weights, batchnorm) {
+             fname_weights, batchnorm, false, final) {
 
     this->kernelH = kernelH;
     this->kernelW = kernelW;

src/DeformConv2d.cpp

@@ -25,25 +25,23 @@ void DeformConv2d::initCUDNN() {
     if (dst_dim % 3 != 0 )
         std::cout<<"take attention\n\n";
     chunk_dim = dst_dim/3;
-    checkCuda(cudaMalloc(&offset, 2*chunk_dim*sizeof(dnnType)));
-    checkCuda(cudaMalloc(&mask, chunk_dim*sizeof(dnnType)));
+    checkCuda( cudaMalloc(&offset, 2*chunk_dim*sizeof(dnnType)));
+    checkCuda( cudaMalloc(&mask, chunk_dim*sizeof(dnnType)));
     
     // kernel ones
     
-    cudaMallocHost(&ones_d1, (height_ones*width_ones)*sizeof(dnnType));
+    checkCuda( cudaMalloc(&ones_d1, (height_ones*width_ones)*sizeof(dnnType)) );
     float aus1[height_ones*width_ones];
     for(int i=0; i<height_ones*width_ones; i++)
         aus1[i]=1.0f;
-    cudaMemcpy(ones_d1, aus1, (height_ones*width_ones)*sizeof(dnnType), cudaMemcpyHostToDevice);
-    cudaDeviceSynchronize();
+    checkCuda( cudaMemcpy(ones_d1, aus1, (height_ones*width_ones)*sizeof(dnnType), cudaMemcpyHostToDevice) );
     
-    cudaMallocHost(&ones_d2, dim_ones*sizeof(dnnType));
+    checkCuda( cudaMalloc(&ones_d2, dim_ones*sizeof(dnnType)) );
     float aus2[dim_ones];
     for(int i=0; i<dim_ones; i++)
         aus2[i]=1.0f;
-    cudaMemcpy(ones_d2, aus2, (dim_ones)*sizeof(dnnType), cudaMemcpyHostToDevice);
-    cudaDeviceSynchronize();
-
+    checkCuda( cudaMemcpy(ones_d2, aus2, (dim_ones)*sizeof(dnnType), cudaMemcpyHostToDevice) );
+    checkCuda( cudaDeviceSynchronize() );
 }
 
 DeformConv2d::DeformConv2d( Network *net, int out_ch, int deformable_group, int kernelH, int kernelW,

src/Layer.cpp

@@ -4,10 +4,10 @@
 
 namespace tk { namespace dnn {
 
-Layer::Layer(Network *net) {
+Layer::Layer(Network *net, bool final) {
 
     this->net = net;
-
+    this->final = final;
     if(net != nullptr) {
         this->input_dim = net->getOutputDim();
         this->output_dim = input_dim;

src/LayerWgs.cpp

@@ -8,12 +8,12 @@ namespace tk { namespace dnn {
 
 LayerWgs::LayerWgs(Network *net, int inputs, int outputs, 
                    int kh, int kw, int kl, 
-                   std::string fname_weights, bool batchnorm, bool additional_bias) : Layer(net) {
+                   std::string fname_weights, bool batchnorm, bool additional_bias, bool final) : Layer(net, final) {
 
     this->inputs  = inputs;
     this->outputs = outputs;    
     this->weights_path  = std::string(fname_weights);
-   
+
     std::cout<<"Reading weights: I="<<inputs<<" O="<<outputs<<" KERNEL="<<kh<<"x"<<kw<<"x"<<kl<<"\n";
     int seek = 0;
     readBinaryFile(weights_path.c_str(), inputs*outputs*kh*kw*kl, &data_h, &data_d, seek, net->dontLoadWeights);

src/NetworkRT.cpp

@@ -75,7 +75,7 @@ NetworkRT::NetworkRT(Network *net, const char *name) {
             input = Ilay->getOutput(0);
             input->setName( (l->getLayerName() + std::to_string(i) + "_out").c_str() );
             
-            if(l->getLayerType() == LAYER_YOLO)
+            if(l->getLayerType() == LAYER_YOLO || l->final)
                 networkRT->markOutput(*input);
             tensors[l] = input;
         }
@@ -182,6 +182,8 @@ ILayer* NetworkRT::convert_layer(ITensor *input, Layer *l) {
         return convert_layer(input, (Yolo*) l);
     if(type == LAYER_UPSAMPLE)
         return convert_layer(input, (Upsample*) l);
+    if(type == LAYER_DEFORMCONV2D)
+        return convert_layer(input, (DeformConv2d*) l);
 
     std::cout<<l->getLayerName()<<"\n";
     FatalError("Layer not implemented in tensorRT");
@@ -254,6 +256,8 @@ ILayer* NetworkRT::convert_layer(ITensor *input, Conv2d *l) {
         lRTconv->setPadding(DimsHW{l->paddingH, l->paddingW});
         lRT = (ILayer*) lRTconv;
         
+        Dims d = lRTconv->getOutput(0)->getDimensions();
+        std::cout<<"DECONV: "<<d.d[0]<<" "<<d.d[1]<<" "<<d.d[2]<<" "<<d.d[3]<<"\n";
     }
 
     checkNULL(lRT);
@@ -423,6 +427,53 @@ ILayer* NetworkRT::convert_layer(ITensor *input, Upsample *l) {
     return lRT;
 }
 
+ILayer* NetworkRT::convert_layer(ITensor *input, DeformConv2d *l) {
+    //std::cout<<"convert DEFORMABLE\n";
+    ILayer *preconv = convert_layer(input, l->preconv);
+
+    ITensor **inputs = new ITensor*[2];
+    inputs[0] = input;
+    inputs[1] = preconv->getOutput(0);
+
+    //std::cout<<"New plugin DEFORMABLE\n";
+    IPlugin *plugin = new DeformableConvRT(l);
+    IPluginLayer *lRT = networkRT->addPlugin(&input, 1, *plugin);
+    checkNULL(lRT);
+
+    // batchnorm
+    void *bias_b, *power_b, *mean_b, *variance_b, *scales_b;
+    if(dtRT == DataType::kHALF) {
+        bias_b     = l->bias16_h;
+        power_b    = l->power16_h;
+        mean_b     = l->mean16_h;
+        variance_b = l->variance16_h;
+        scales_b   = l->scales16_h;
+    } else {
+        bias_b     = l->bias_h;
+        power_b    = l->power_h;
+        mean_b     = l->mean_h;
+        variance_b = l->variance_h;
+        scales_b   = l->scales_h;
+    }
+
+    Weights power{dtRT, power_b, l->outputs};
+    Weights shift{dtRT, mean_b, l->outputs};
+    Weights scale{dtRT, variance_b, l->outputs};
+    std::cout<<lRT->getNbOutputs()<<std::endl;
+    IScaleLayer *lRT2 = networkRT->addScale(*lRT->getOutput(0), ScaleMode::kCHANNEL, 
+                shift, scale, power);
+    
+    checkNULL(lRT2);
+
+    Weights shift2{dtRT, bias_b, l->outputs};
+    Weights scale2{dtRT, scales_b, l->outputs};
+    IScaleLayer *lRT3 = networkRT->addScale(*lRT2->getOutput(0), ScaleMode::kCHANNEL, 
+                shift2, scale2, power);
+    checkNULL(lRT3);
+
+    return lRT3;
+}
+
 bool NetworkRT::serialize(const char *filename) {
 
     std::ofstream p(filename);