findCUDNN

cmake/FindCUDNN.cmake

@@ -0,0 +1,33 @@
+# Find the header files
+
+find_path(CUDNN_INCLUDE_DIR
+        ${CMAKE_SYSROOT}/usr/local/include
+        ${CMAKE_SYSROOT}/usr/include
+        /usr/local/nvidia/tensorrt/include/
+        NO_DEFAULT_PATH
+)
+
+set(OLD_ROOT ${CMAKE_FIND_ROOT_PATH})
+list(APPEND CMAKE_FIND_ROOT_PATH /)
+list(APPEND CMAKE_FIND_LIBRARY_SUFFIXES .so.7)
+list(APPEND CMAKE_FIND_LIBRARY_SUFFIXES .so.5)
+find_library(CUDNN_LIB
+    NAMES cudnn
+    PATHS
+    /usr/local/driveworks-2.0/targets/${CMAKE_SYSTEM_PROCESSOR}-Linux/lib
+    /usr/lib/${CMAKE_SYSTEM_PROCESSOR}-linux-gnu/
+    NO_DEFAULT_PATH
+)
+find_library(CUDNN_NVLIB
+    NAMES "nvinfer"
+    PATHS
+    /usr/local/driveworks-2.0/targets/${CMAKE_SYSTEM_PROCESSOR}-Linux/lib
+    /usr/lib/${CMAKE_SYSTEM_PROCESSOR}-linux-gnu/
+    NO_DEFAULT_PATH
+)
+set(CMAKE_FIND_ROOT_PATH ${OLD_ROOT})
+
+set(CUDNN_LIBRARIES ${CUDNN_LIB} ${CUDNN_NVLIB})
+message("-- Found CUDNN: "  ${CUDNN_LIB})
+message("-- Found NVINFER: "  ${CUDNN_NVLIB})
+set(CUDNN_FOUND true)

cmake/tkDNNConfig.cmake

@@ -4,27 +4,21 @@ set(CMAKE_CXX_FLAGS  "${CMAKE_CXX_FLAGS} --std=c++11 -fPIC")
 
 find_package(CUDA REQUIRED)
 find_package(OpenCV REQUIRED)
-find_library(NVINFER NAMES nvinfer)
-if(NVINFER STREQUAL "NVINFER-NOTFOUND")
-    set(NVINFER_INCLUDES "/usr/local/nvidia/tensorrt/include/")
-    link_directories(/usr/local/nvidia/tensorrt/targets/x86_64-linux-gnu/lib/ 
-		     /usr/local/cuda/targets/x86_64-linux/lib/)
-endif()
+find_package(CUDNN REQUIRED)
 
 set(tkDNN_INCLUDE_DIRS 
 	${CUDA_INCLUDE_DIRS} 
 	${OPENCV_INCLUDE_DIRS} 
-    ${NVINFER_INCLUDES}
+    ${CUDNN_INCLUDE_DIRS}
 )
 
 set(tkDNN_LIBRARIES 
     tkDNN 
     kernels 
     ${CUDA_LIBRARIES} 
-    ${CUDA_CUBLAS_LIBRARIES} 
-    -lcudnn 
-    -lnvinfer 
-    ${OpenCV_LIBS}
+    ${CUDA_CUBLAS_LIBRARIES}
+	${CUDNN_LIBRARIES}
+	${OpenCV_LIBS}
 )
 
 set(tkDNN_FOUND true)