TensorRT version

src/Conv2d.cpp

@@ -37,7 +37,7 @@ Conv2d::Conv2d( Network *net, dataDim_t in_dim, int out_ch,
                 paddingH, paddingW, // padding
                 strideH, strideW, // stride
                 1,1, // upscale
-                CUDNN_CROSS_CORRELATION) );
+                CUDNN_CROSS_CORRELATION, CUDNN_DATA_FLOAT) );
 
     // find dimension of convolution output
     checkCUDNN( cudnnGetConvolution2dForwardOutputDim(

src/Network.cpp

@@ -1,4 +1,5 @@
 #include <iostream>
+#include "NvInfer.h"
 
 #include "tkdnn.h"
 #include "Network.h"
@@ -10,8 +11,10 @@ Network::Network() {
 
     float tk_ver = float(tkDNN::getVersion())/1000;
     float cu_ver = float(cudnnGetVersion())/1000;
+    float rt_ver = float(NV_TENSORRT_MAJOR) + float(NV_TENSORRT_MINOR)/10 + float(NV_TENSORRT_PATCH)/100;
 
-    std::cout<<"New NETWORK (tkDNN v"<<tk_ver<<", CUDNN v"<<cu_ver<<")\n";
+    std::cout<<"New NETWORK (tkDNN v"<<tk_ver
+             <<", CUDNN v"<<cu_ver<<", TensorRT v"<<rt_ver<<")\n";
     dataType = CUDNN_DATA_FLOAT;
     tensorFormat = CUDNN_TENSOR_NCHW;
 
@@ -32,7 +35,7 @@ value_type* Network::infer(dataDim_t &dim, value_type* data) {
     //do infer for every layer
     for(int i=0; i<num_layers; i++)
         data = layers[i]->infer(dim, data);
-        
+
     return data;
 }
 

src/Reorg.cpp

@@ -26,7 +26,7 @@ Reorg::~Reorg() {
 
 value_type* Reorg::infer(dataDim_t &dim, value_type* srcData) {
 
-    reorgForward(srcData, dstData, dim, stride);
+    reorgForward(srcData, dstData, dim.n, dim.c, dim.h, dim.w, stride);
 
     dim = output_dim;
     return dstData;

src/kernels/reorg.cu

@@ -35,14 +35,15 @@ __global__ void reorg_kernel(int N, float *x, int w, int h, int c, int batch, in
 /**
     reorg function function
 */
-void reorgForward(value_type* srcData, value_type* dstData, tkDNN::dataDim_t dim, int stride)
-{
-    int size = dim.tot();
+void reorgForward(value_type* srcData, value_type* dstData, 
+                  int n, int c, int h, int w, int stride) {
+
+    int size = n*c*h*w;
 
     int blocks = (size+255)/256;
     int threads = 256;
     
-    reorg_kernel<<<blocks, threads>>>(size, srcData, dim.w, dim.h, dim.c, dim.n, stride, false, dstData);
+    reorg_kernel<<<blocks, threads>>>(size, srcData, w, h, c, n, stride, false, dstData);
     checkCuda( cudaDeviceSynchronize() );
 }