RTinference test

CMakeLists.txt

@@ -54,3 +54,6 @@ target_link_libraries(test_yolo tkDNN)
 
 add_executable(test_yolo_tiny tests/yolo-tiny/yolo-tiny.cpp)
 target_link_libraries(test_yolo_tiny tkDNN)
+
+add_executable(test_rtinference tests/test_rtinference/rtinference.cpp)
+target_link_libraries(test_rtinference tkDNN)

include/NetworkRT.h

@@ -21,7 +21,7 @@ public:
     void* buffersRT[2];
     int buf_input_idx, buf_output_idx;
 
-    dataDim_t output_dim;
+    dataDim_t input_dim, output_dim;
     dnnType *output;
     cudaStream_t stream;
 

src/NetworkRT.cpp

@@ -36,9 +36,10 @@ NetworkRT::NetworkRT(Network *net, const char *name) {
     networkRT = builderRT->createNetwork();
     dtRT = DataType::kFLOAT;
 
-    //add input layer
-    dataDim_t dim = net->layers[0]->input_dim;
     if(!fileExist(name)) {
+        //add input layer
+        dataDim_t dim = net->layers[0]->input_dim;
+        
         ITensor *input = networkRT->addInput("data", dtRT, 
                         DimsCHW{ dim.c, dim.h, dim.w});
         checkNULL(input);
@@ -91,14 +92,21 @@ NetworkRT::NetworkRT(Network *net, const char *name) {
     buf_output_idx = engineRT->getBindingIndex("out");
     std::cout<<"input idex = "<<buf_input_idx<<" -> output index = "<<buf_output_idx<<"\n";
 
-    output_dim = dim;
+
+    Dims iDim = engineRT->getBindingDimensions(buf_output_idx);
+    input_dim.n = 1;
+    input_dim.c = iDim.d[0];
+    input_dim.h = iDim.d[1];
+    input_dim.w = iDim.d[2];
+
     Dims oDim = engineRT->getBindingDimensions(buf_output_idx);
+    output_dim.n = 1;
     output_dim.c = oDim.d[0];
     output_dim.h = oDim.d[1];
     output_dim.w = oDim.d[2];
 	
     // create GPU buffers and a stream
-    checkCuda(cudaMalloc(&buffersRT[buf_input_idx],  dim.tot()*sizeof(dnnType)));
+    checkCuda(cudaMalloc(&buffersRT[buf_input_idx],  input_dim.tot()*sizeof(dnnType)));
     checkCuda(cudaMalloc(&buffersRT[buf_output_idx], output_dim.tot()*sizeof(dnnType)));
     checkCuda(cudaMalloc(&output, output_dim.tot()*sizeof(dnnType)));
 	checkCuda(cudaStreamCreate(&stream));
@@ -110,7 +118,7 @@ NetworkRT::~NetworkRT() {
 
 dnnType* NetworkRT::infer(dataDim_t &dim, dnnType* data) {
 
-    checkCuda(cudaMemcpyAsync(buffersRT[buf_input_idx], data, dim.tot()*sizeof(float), cudaMemcpyDeviceToDevice, stream));
+    checkCuda(cudaMemcpyAsync(buffersRT[buf_input_idx], data, input_dim.tot()*sizeof(float), cudaMemcpyDeviceToDevice, stream));
     contextRT->enqueue(1, buffersRT, stream, nullptr);
     checkCuda(cudaMemcpyAsync(output, buffersRT[buf_output_idx], output_dim.tot()*sizeof(float), cudaMemcpyDeviceToDevice, stream));
     cudaStreamSynchronize(stream);

tests/test_rtinference/rtinference.cpp

@@ -0,0 +1,26 @@
+#include<iostream>
+#include "tkdnn.h"
+
+int main(int argc, char *argv[]) {
+
+    // Network layout
+    tkDNN::dataDim_t dim(1, 3, 608, 608, 1);
+    tkDNN::Network net(dim);
+
+    if(argc < 2 || !fileExist(argv[1]))
+        FatalError("unable to read serialRT file");
+
+    //convert network to tensorRT
+    tkDNN::NetworkRT netRT(&net, argv[1]);
+
+    dnnType *data;
+    checkCuda(cudaMalloc(&data, dim.tot()*sizeof(dnnType)));
+
+    printCenteredTitle(" TENSORRT inference ", '=', 30); {
+        TIMER_START
+        data = netRT.infer(dim, data);
+        TIMER_STOP
+    }
+    
+    return 0;
+}

tests/yolo/yolo.cpp

@@ -28,7 +28,7 @@ const char *c30_bin    = "../tests/yolo/layers/c30.bin";
 const char *g31_bin    = "../tests/yolo/layers/g31.bin";
 const char *output_bin = "../tests/yolo/layers/output.bin";
 
-int main(int argc, char *argv[]) {
+int main() {
 
     // Network layout
     tkDNN::dataDim_t dim(1, 3, 608, 608, 1);
@@ -108,7 +108,7 @@ int main(int argc, char *argv[]) {
     net.print();
 
     //convert network to tensorRT
-    tkDNN::NetworkRT netRT(&net, argv[1]);
+    tkDNN::NetworkRT netRT(&net, "yolo.rt");
 
     dnnType *out_data, *out_data2; // cudnn output, tensorRT output