From b3a369dc293ee268627faca927d9124d35ef95b7 Mon Sep 17 00:00:00 2001
From: Francesco Gatti <gattifrancesco@hotmail.it>
Date: Mon, 14 Aug 2017 11:24:23 +0200
Subject: [PATCH] RTinference test

---
 CMakeLists.txt                         |  3 +++
 include/NetworkRT.h                    |  2 +-
 src/NetworkRT.cpp                      | 18 +++++++++++++-----
 tests/test_rtinference/rtinference.cpp | 26 ++++++++++++++++++++++++++
 tests/yolo/yolo.cpp                    |  4 ++--
 5 files changed, 45 insertions(+), 8 deletions(-)
 create mode 100644 tests/test_rtinference/rtinference.cpp

diff --git a/CMakeLists.txt b/CMakeLists.txt
index afc523c..a1436bc 100644
--- a/CMakeLists.txt
+++ b/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)
\ No newline at end of file
diff --git a/include/NetworkRT.h b/include/NetworkRT.h
index 0a57e2a..710c4d4 100644
--- a/include/NetworkRT.h
+++ b/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;
 
diff --git a/src/NetworkRT.cpp b/src/NetworkRT.cpp
index b891377..26fa1d2 100644
--- a/src/NetworkRT.cpp
+++ b/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);
diff --git a/tests/test_rtinference/rtinference.cpp b/tests/test_rtinference/rtinference.cpp
new file mode 100644
index 0000000..3c61083
--- /dev/null
+++ b/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;
+}
diff --git a/tests/yolo/yolo.cpp b/tests/yolo/yolo.cpp
index 7c08a32..50aec86 100644
--- a/tests/yolo/yolo.cpp
+++ b/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