better print

CMakeLists.txt

@@ -13,6 +13,10 @@ if( ${BUILD_DEPS} )
 	message("Finished dowloading test weights")
 endif()
 
+if(DEBUG)
+    add_definitions(-DDEBUG)
+endif()
+
 find_package(CUDA QUIET REQUIRED)
 
 cuda_include_directories(${CMAKE_CURRENT_SOURCE_DIR}/include ${CUDA_INCLUDE_DIRS})

src/Network.cpp

@@ -58,7 +58,6 @@ dataDim_t Network::getOutputDim() {
 
 void Network::print() {
 
-    std::cout<<"\n";
     printCenteredTitle(" NETWORK MODEL ", '=', 60);
     std::cout.width(3); std::cout<<std::left<<"N.";
     std::cout<<" ";

src/NetworkRT.cpp

@@ -12,7 +12,9 @@ using namespace nvinfer1;
 // Logger for info/warning/errors
 class Logger : public ILogger {
     void log(Severity severity, const char* msg) override {
+#ifdef DEBUG
         std::cout <<"TENSORRT LOG: "<< msg << std::endl;
+#endif
     }
 } loggerRT;
 
@@ -59,7 +61,7 @@ NetworkRT::NetworkRT(Network *net) {
 	builderRT->setMaxBatchSize(1);
 	builderRT->setMaxWorkspaceSize(1 << 20);
 
-    std::cout<<"BUILD cuda engine\n";
+    std::cout<<"Building tensorRT cuda engine...\n";
 	engineRT = builderRT->buildCudaEngine(*networkRT);
 	// we don't need the network any more
 	//networkRT->destroy();
@@ -127,7 +129,7 @@ ITensor* NetworkRT::convert_layer(ITensor *input, Layer *l) {
 }
 
 ITensor* NetworkRT::convert_layer(ITensor *input, Dense *l) {
-    std::cout<<"convert Dense\n";
+    //std::cout<<"convert Dense\n";
 
     Weights w { dtRT, l->data_h, l->inputs*l->outputs};
     Weights b = { dtRT, l->bias_h, l->outputs};
@@ -138,7 +140,7 @@ ITensor* NetworkRT::convert_layer(ITensor *input, Dense *l) {
 }
 
 ITensor* NetworkRT::convert_layer(ITensor *input, Conv2d *l) {
-    std::cout<<"convert conv2D\n";
+    //std::cout<<"convert conv2D\n";
 
     Weights w { dtRT, l->data_h, l->inputs*l->outputs*l->kernelH*l->kernelW};
     Weights b;
@@ -190,7 +192,7 @@ ITensor* NetworkRT::convert_layer(ITensor *input, Conv2d *l) {
 }
 
 ITensor* NetworkRT::convert_layer(ITensor *input, Pooling *l) {
-    std::cout<<"convert Pooling\n";
+    //std::cout<<"convert Pooling\n";
 
     IPoolingLayer *lRT = networkRT->addPooling(*input, 
         PoolingType::kMAX, DimsHW{l->winH, l->winW});
@@ -201,10 +203,10 @@ ITensor* NetworkRT::convert_layer(ITensor *input, Pooling *l) {
 }
 
 ITensor* NetworkRT::convert_layer(ITensor *input, Activation *l) {
-    std::cout<<"convert Activation\n";
+    //std::cout<<"convert Activation\n";
 
     if(l->act_mode == ACTIVATION_LEAKY) {
-        std::cout<<"New plugin LEAKY\n";
+        //std::cout<<"New plugin LEAKY\n";
         IPlugin *plugin = new ActivationLeakyRT();
         IPluginLayer *lRT = networkRT->addPlugin(&input, 1, *plugin);
         checkNULL(lRT);
@@ -217,7 +219,7 @@ ITensor* NetworkRT::convert_layer(ITensor *input, Activation *l) {
 }
 
 ITensor* NetworkRT::convert_layer(ITensor *input, Softmax *l) {
-    std::cout<<"convert softmax\n";
+    //std::cout<<"convert softmax\n";
 
     ISoftMaxLayer *lRT = networkRT->addSoftMax(*input);
     checkNULL(lRT);
@@ -226,7 +228,7 @@ ITensor* NetworkRT::convert_layer(ITensor *input, Softmax *l) {
 }
 
 ITensor* NetworkRT::convert_layer(ITensor *input, Route *l) {
-    std::cout<<"convert route\n";
+    //std::cout<<"convert route\n";
 
     ITensor *tens[256];
     for(int i=0; i<l->layers_n; i++)
@@ -238,9 +240,9 @@ ITensor* NetworkRT::convert_layer(ITensor *input, Route *l) {
 }
 
 ITensor* NetworkRT::convert_layer(ITensor *input, Reorg *l) {
-    std::cout<<"convert Reorg\n";
+    //std::cout<<"convert Reorg\n";
 
-    std::cout<<"New plugin REORG\n";
+    //std::cout<<"New plugin REORG\n";
     IPlugin *plugin = new ReorgRT(l->stride);
     IPluginLayer *lRT = networkRT->addPlugin(&input, 1, *plugin);
     checkNULL(lRT);
@@ -248,9 +250,9 @@ ITensor* NetworkRT::convert_layer(ITensor *input, Reorg *l) {
 }
 
 ITensor* NetworkRT::convert_layer(ITensor *input, Region *l) {
-    std::cout<<"convert Region\n";
+    //std::cout<<"convert Region\n";
 
-    std::cout<<"New plugin REGION\n";
+    //std::cout<<"New plugin REGION\n";
     IPlugin *plugin = new RegionRT(l->classes, l->coords, l->num, l->thresh);
     IPluginLayer *lRT = networkRT->addPlugin(&input, 1, *plugin);
     checkNULL(lRT);

src/utils.cpp

@@ -6,6 +6,8 @@ void printCenteredTitle(const char *title, char fill, int dim) {
     int len = strlen(title);
     int first = dim/2 + len/2;
     
+    if(len >0)
+        std::cout<<"\n";
     std::cout.width(first); std::cout.fill(fill); std::cout<<std::right<<title;
     std::cout.width(dim - first); std::cout<<"\n"; 
     std::cout.fill(' ');

tests/yolo-tiny/yolo-tiny.cpp

@@ -56,13 +56,16 @@ int main() {
     value_type *input_h;
     readBinaryFile(input_bin, dim.tot(), &input_h, &data);
 
+    //print network model
+    net.print();
+
     //convert network to tensorRT
     tkDNN::NetworkRT netRT(&net);
 
     value_type *out_data, *out_data2; // cudnn output, tensorRT output
 
     tkDNN::dataDim_t dim1 = dim; //input dim
-    std::cout<<"\n==== CUDNN inference =======\n"; {
+    printCenteredTitle(" CUDNN inference ", '=', 30); {
         dim1.print();
         TIMER_START
         out_data = net.infer(dim1, data);    
@@ -71,7 +74,7 @@ int main() {
     }
  
     tkDNN::dataDim_t dim2 = dim;
-    std::cout<<"\n==== TENSORRT inference ====\n"; {
+    printCenteredTitle(" TENSORRT inference ", '=', 30); {
         dim2.print();
         TIMER_START
         out_data2 = netRT.infer(dim2, data);
@@ -79,7 +82,7 @@ int main() {
         dim2.print();
     }
 
-    std::cout<<"\n======= CHECK RESULT =======\n";
+    printCenteredTitle(" CHECK RESULTS ", '=', 30);
     value_type *out, *out_h;
     int out_dim = net.getOutputDim().tot();
     readBinaryFile(output_bin, out_dim, &out_h, &out);

tests/yolo/yolo.cpp

@@ -104,13 +104,16 @@ int main() {
     value_type *input_h;
     readBinaryFile(input_bin, dim.tot(), &input_h, &data);
     
+    //print network model
+    net.print();
+
     //convert network to tensorRT
     tkDNN::NetworkRT netRT(&net);
 
     value_type *out_data, *out_data2; // cudnn output, tensorRT output
 
     tkDNN::dataDim_t dim1 = dim; //input dim
-    std::cout<<"\n==== CUDNN inference =======\n"; {
+    printCenteredTitle(" CUDNN inference ", '=', 30); {
         dim1.print();
         TIMER_START
         out_data = net.infer(dim1, data);    
@@ -119,7 +122,7 @@ int main() {
     }
  
     tkDNN::dataDim_t dim2 = dim;
-    std::cout<<"\n==== TENSORRT inference ====\n"; {
+    printCenteredTitle(" TENSORRT inference ", '=', 30); {
         dim2.print();
         TIMER_START
         out_data2 = netRT.infer(dim2, data);
@@ -127,7 +130,7 @@ int main() {
         dim2.print();
     }
 
-    std::cout<<"\n======= CHECK RESULT =======\n";
+    printCenteredTitle(" CHECK RESULTS ", '=', 30);
     value_type *out, *out_h;
     int out_dim = net.getOutputDim().tot();
     readBinaryFile(output_bin, out_dim, &out_h, &out);