CMakeLists.txt

@@ -6,6 +6,32 @@ set(CMAKE_CXX_STANDARD 14)
 option(ENABLE_OPENCV_CUDA_CONTRIB "Enable OpenCV CUDA Contrib" OFF )
 
 find_package(CUDA 9.0 REQUIRED)
+if (CUDA_FOUND)
+    set(OUTPUTFILE ${CMAKE_CURRENT_SOURCE_DIR}/cmake/cuda_script) # No suffix required
+    execute_process(COMMAND "rm ${OUTPUTFILE}")
+    set(CUDAFILE ${CMAKE_CURRENT_SOURCE_DIR}/cmake/getCudaArch.cu)
+    execute_process(COMMAND ${CUDA_NVCC_EXECUTABLE} -lcuda ${CUDAFILE} -o ${OUTPUTFILE})
+    execute_process(COMMAND ${OUTPUTFILE}
+            RESULT_VARIABLE CUDA_RETURN_CODE
+            OUTPUT_VARIABLE ARCH)
+
+    if(${CUDA_RETURN_CODE} EQUAL 0)
+        set(CUDA_SUCCESS "TRUE")
+    else()
+        set(CUDA_SUCCESS "FALSE")
+    endif()
+
+    if (${CUDA_SUCCESS})
+        message(STATUS "CUDA Architecture: ${ARCH}")
+        message(STATUS "CUDA Version: ${CUDA_VERSION_STRING}")
+        message(STATUS "CUDA Path: ${CUDA_TOOLKIT_ROOT_DIR}")
+        message(STATUS "CUDA Libararies: ${CUDA_LIBRARIES}")
+        message(STATUS "CUDA Performance Primitives: ${CUDA_npp_LIBRARY}")
+        set(CUDA_NVCC_FLAGS "${ARCH}")
+    else()
+        message(WARNING ${ARCH})
+    endif()
+endif()
 
 SET(CUDA_SEPARABLE_COMPILATION ON)
 

include/tkDNN/NetworkRT.h

@@ -24,6 +24,8 @@
 #include <pluginsRT/UpsampleRT.h>
 #include <pluginsRT/YoloRT.h>
 
+
+
 namespace tk { namespace dnn {
 
 using namespace nvinfer1;

src/NetworkRT.cpp

@@ -183,7 +183,11 @@ NetworkRT::NetworkRT(Network *net, const char *name) {
     output_dim.h = oDim.d[1];
     output_dim.w = oDim.d[2];
     output_dim.print();
-	
+    if(builderActive){
+        std::cout<<"NUMBER OF LAYERS IN NETWORK : "<<networkRT->getNbLayers()<<std::endl;
+    }
+    std::cout<<"NUMBER OF LAYERS IN ENGINE : "<<engineRT->getNbLayers()<<std::endl;
+
     // create GPU buffers and a stream
     for(int i=0; i<engineRT->getNbBindings(); i++) {
         Dims dim = engineRT->getBindingDimensions(i);