findCUDNN

2019-11-06 14:04:23 +01:00
parent 92f3d1c548
commit aa5927d8a1
3 changed files with 41 additions and 20 deletions
@@ -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)