fix

include/utils.h

@@ -14,14 +14,31 @@
 
 #define value_type float
 
+// Colored output
+#define COL_END "\033[0m"
 
+#define COL_RED "\033[31m"
+#define COL_GREEN "\033[32m"
+#define COL_ORANGE "\033[33m"
+#define COL_BLUE "\033[34m"
+#define COL_PURPLE "\033[35m"
+#define COL_CYAN "\033[36m"
+
+#define COL_REDB "\033[1;31m"
+#define COL_GREENB "\033[1;32m"
+#define COL_ORANGEB "\033[1;33m"
+#define COL_BLUEB "\033[1;34m"
+#define COL_PURPLEB "\033[1;35m"
+#define COL_CYANB "\033[1;36m"
+
+// Simple Timer 
 #define TIMER_START timespec start, end;                               \
                     clock_gettime(CLOCK_MONOTONIC, &start);            
 
 #define TIMER_STOP  clock_gettime(CLOCK_MONOTONIC, &end);              \
     double t_ns = ((double)(end.tv_sec - start.tv_sec) * 1.0e9 +       \
                   (double)(end.tv_nsec - start.tv_nsec))/1.0e6;        \
-    std::cout<<"\033[1;36mTime:"<<std::setw(16)<<t_ns<<" ms\033[0m\n"; 
+    std::cout<<COL_CYANB<<"Time:"<<std::setw(16)<<t_ns<<" ms\n"<<COL_END; 
 
 
 /********************************************************
@@ -71,7 +88,7 @@
 }
 
 void readBinaryFile(const char* fname, int size, value_type** data_h, value_type** data_d, int seek = 0);
-int  checkResult(int size, value_type *data_d, value_type *correct_d, bool device = true);
+int checkResult(int size, value_type *data_d, value_type *correct_d, bool device = true);
 void printDeviceVector(int size, value_type* vec_d, bool device = true);
 void resize(int size, value_type **data);
 

src/Network.cpp

@@ -1,5 +1,4 @@
 #include <iostream>
-#include "NvInfer.h"
 
 #include "tkdnn.h"
 #include "Network.h"
@@ -12,10 +11,9 @@ Network::Network(dataDim_t input_dim) {
 
     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<<", TensorRT v"<<rt_ver<<")\n";
+             <<", CUDNN v"<<cu_ver<<")\n";
     dataType = CUDNN_DATA_FLOAT;
     tensorFormat = CUDNN_TENSOR_NCHW;
 

src/NetworkRT.cpp

@@ -10,12 +10,10 @@ using namespace nvinfer1;
 #include "pluginsRT/RegionRT.cpp"
 
 // Logger for info/warning/errors
-class Logger : public ILogger			
+class Logger : public ILogger {
-{
+    void log(Severity severity, const char* msg) override {
-	void log(Severity severity, const char* msg) override
+        std::cout <<"TENSORRT LOG: "<< msg << std::endl;
-	{
+    }
-		std::cout <<"TENSORRT: "<< msg << std::endl;
-	}
 } loggerRT;
 
 namespace tkDNN {
@@ -24,6 +22,11 @@ std::map<Layer*, nvinfer1::ITensor*>tensors;
 
 NetworkRT::NetworkRT(Network *net) {
 
+    float rt_ver = float(NV_TENSORRT_MAJOR) + 
+                   float(NV_TENSORRT_MINOR)/10 + 
+                   float(NV_TENSORRT_PATCH)/100;
+    std::cout<<"New NetworkRT (TensorRT v"<<rt_ver<<")\n";
+
     builderRT = createInferBuilder(loggerRT);
     networkRT = builderRT->createNetwork();
     dtRT = DataType::kFLOAT;

src/utils.cpp

@@ -23,9 +23,7 @@ void readBinaryFile(const char* fname, int size, value_type** data_h, value_type
     }
     
     checkCuda( cudaMalloc(data_d, size_b) );
-    checkCuda( cudaMemcpy(*data_d, *data_h,
+    checkCuda( cudaMemcpy(*data_d, *data_h, size_b, cudaMemcpyHostToDevice) );
-                                size_b,
-                                cudaMemcpyHostToDevice) );
 }
 
 void printDeviceVector(int size, value_type* vec_d, bool device)
@@ -39,8 +37,7 @@ void printDeviceVector(int size, value_type* vec_d, bool device)
         vec = vec_d;
     }
     
-    for (int i = 0; i < size; i++)
+    for (int i = 0; i < size; i++) {
-    {
         std::cout << vec[i] << " ";
     }
     std::cout << std::endl;
@@ -52,6 +49,7 @@ void printDeviceVector(int size, value_type* vec_d, bool device)
 int checkResult(int size, value_type *data_d, value_type *correct_d, bool device) {
 
     value_type *data_h, *correct_h;
+    const float eps = 0.0001f;
 
     if(device) {
         data_h = new value_type[size];
@@ -67,10 +65,12 @@ int checkResult(int size, value_type *data_d, value_type *correct_d, bool device
 
     int diffs = 0;
     for(int i=0; i<size; i++) {
-        if(fabs(data_h[i] - correct_h[i]) > 0.0001) {
+        if(fabs(data_h[i] - correct_h[i]) > eps) {
             diffs += 1;
+            if(diffs == 1)
+                std::cout<<"\n";
             if(diffs < 10)
-                printf("%d: %f %f\n", i, data_h[i], correct_h[i]);
+                std::cout<<" | [ "<<i<<" ]: "<<data_h[i]<<" "<<correct_h[i]<<"\n";
         }
     }
 
@@ -79,6 +79,13 @@ int checkResult(int size, value_type *data_d, value_type *correct_d, bool device
         delete [] correct_h;
     }
 
+    std::cout<<" | ";
+    if(diffs == 0)
+        std::cout<<COL_GREENB<<"OK";
+    else
+        std::cout<<COL_REDB<<"Wrongs: "<<diffs;
+
+    std::cout<<COL_END<<" ~"<<eps<<"\n";
     return diffs;
 }
 

tests/mnist/test_mnist.cpp

@@ -58,7 +58,7 @@ int main() {
     //printDeviceVector(10, out_data);
 
     std::cout<<"\n======= CHECK RESULT =======\n";
-    std::cout<<"Diffs: "<<checkResult(dim.tot(), out_data, out_data2)<<"\n";
+    checkResult(dim.tot(), out_data, out_data2);
 
  /*
     // Print real test

tests/yolo/yolo.cpp

@@ -103,13 +103,14 @@ int main() {
     value_type *input_h;
     readBinaryFile(input_bin, dim.tot(), &input_h, &data);
 
+    //convert network to tensorRT
     tkDNN::NetworkRT netRT(&net);
 
-    value_type *out_data, *out_data2;
+    value_type *out_data, *out_data2; // cudnn output, tensorRT output
 
-    tkDNN::dataDim_t dim1 = dim;
+    tkDNN::dataDim_t dim1 = dim; //input dim
     std::cout<<"\n==== CUDNN inference =======\n"; {
-        dim1.print(); //print initial dimension  
+        dim1.print();
         TIMER_START
         out_data = net.infer(dim1, data);    
         TIMER_STOP
@@ -129,8 +130,8 @@ int main() {
     value_type *out, *out_h;
     int out_dim = net.getOutputDim().tot();
     readBinaryFile(output_bin, out_dim, &out_h, &out);
-    std::cout<<"CUDNN vs correct Wrongs: "<<checkResult(out_dim, out_data, out)<<"\n";
+    std::cout<<"CUDNN vs correct"; checkResult(out_dim, out_data, out);
-    std::cout<<"TRT   vs correct Wrongs: "<<checkResult(out_dim, out_data2, out)<<"\n";
+    std::cout<<"TRT   vs correct"; checkResult(out_dim, out_data2, out);
-    std::cout<<"CUDNN vs TRT     Wrongs: "<<checkResult(out_dim, out_data, out)<<"\n";
+    std::cout<<"CUDNN vs TRT    "; checkResult(out_dim, out_data, out_data2);
     return 0;
 }