Deformable batch works

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

include/tkDNN/kernels.h

@@ -41,7 +41,7 @@ void dcnV2CudaForward(cublasStatus_t stat, cublasHandle_t handle,
                          const int stride_h, const int stride_w,
                          const int pad_h, const int pad_w,
                          const int dilation_h, const int dilation_w,
-                         const int deformable_group,
+                         const int deformable_group, const int batch_id,
                          const int in_n, const int in_c, const int in_h, const int in_w,
                          const int out_n, const int out_c, const int out_h, const int out_w,
                          const int dst_dim, cudaStream_t stream = cudaStream_t(0));

include/tkDNN/pluginsRT/DeformableConvRT.h

@@ -86,26 +86,26 @@ public:
 		dnnType *output_conv = (dnnType*)reinterpret_cast<const dnnType*>(inputs[1]);
 
 		// split conv2d outputs into offset to mask
-		checkCuda(cudaMemcpy(offset, output_conv, 2*chunk_dim*sizeof(dnnType), cudaMemcpyDeviceToDevice)); 
-		checkCuda(cudaMemcpy(mask, output_conv + 2*chunk_dim, chunk_dim*sizeof(dnnType), cudaMemcpyDeviceToDevice)); 
-		// kernel sigmoide
-		activationSIGMOIDForward(mask, mask, chunk_dim);
-	
-		// deformable convolution
-		dcnV2CudaForward(stat, handle, 
-                         	srcData, data_d,
-							bias2_d, ones_d1,
-							offset, mask,
-							reinterpret_cast<dnnType*>(outputs[0]), ones_d2,
-							kh, kw,
-							sh, sw,
-							ph, pw,
-							1, 1,
-							deformableGroup, 
-							i_n, i_c, i_h, i_w,
-							o_n, o_c, o_h, o_w,
-							chunk_dim);
-
+		for(int b=0; b<batchSize; b++) {
+			checkCuda(cudaMemcpy(offset, output_conv + b * 3 * chunk_dim, 2*chunk_dim*sizeof(dnnType), cudaMemcpyDeviceToDevice)); 
+			checkCuda(cudaMemcpy(mask, output_conv + b * 3 * chunk_dim + 2*chunk_dim, chunk_dim*sizeof(dnnType), cudaMemcpyDeviceToDevice)); 
+			// kernel sigmoide
+			activationSIGMOIDForward(mask, mask, chunk_dim);
+			// deformable convolution
+			dcnV2CudaForward(stat, handle, 
+								srcData, data_d,
+								bias2_d, ones_d1,
+								offset, mask,
+								reinterpret_cast<dnnType*>(outputs[0]), ones_d2,
+								kh, kw,
+								sh, sw,
+								ph, pw,
+								1, 1,
+								deformableGroup, b,
+								i_n, i_c, i_h, i_w,
+								o_n, o_c, o_h, o_w,
+								chunk_dim);
+		}
 		return 0;
 	}
 
@@ -185,6 +185,11 @@ public:
 	dnnType * offset;
 	dnnType * mask;
 	dnnType *ones_d2;
+	// dnnType *input_n;
+	// dnnType *offset_n;
+	// dnnType *mask_n;
+	// dnnType *output_n;
+	
 	
 	tk::dnn::DeformConv2d *defRT;
 };