findCUDNN
This commit is contained in:
+3
-9
@@ -18,17 +18,11 @@ find_package(CUDA 9.0 REQUIRED)
|
||||
SET(CUDA_SEPARABLE_COMPILATION ON)
|
||||
#set(CUDA_NVCC_FLAGS "${CUDA_NVCC_FLAGS} -arch=sm_30 --compiler-options '-fPIC'")
|
||||
|
||||
# compile Discovery only if TensorRT is installed
|
||||
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)
|
||||
|
||||
# compile
|
||||
file(GLOB tkdnn_CUSRC "src/kernels/*.cu")
|
||||
cuda_include_directories(${CMAKE_CURRENT_SOURCE_DIR}/include ${CUDA_INCLUDE_DIRS} ${NVINFER_INCLUDES})
|
||||
cuda_include_directories(${CMAKE_CURRENT_SOURCE_DIR}/include ${CUDA_INCLUDE_DIRS} ${CUDNN_INCLUDE_DIRS})
|
||||
cuda_add_library(kernels SHARED ${tkdnn_CUSRC})
|
||||
|
||||
|
||||
@@ -43,7 +37,7 @@ set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -DOPENCV")
|
||||
# Build Libraries
|
||||
#-------------------------------------------------------------------------------
|
||||
file(GLOB tkdnn_SRC "src/*.cpp")
|
||||
set(tkdnn_LIBS kernels ${CUDA_LIBRARIES} ${CUDA_CUBLAS_LIBRARIES} -lcudnn -lnvinfer ${OpenCV_LIBS})
|
||||
set(tkdnn_LIBS kernels ${CUDA_LIBRARIES} ${CUDA_CUBLAS_LIBRARIES} ${CUDNN_LIBRARIES} ${OpenCV_LIBS})
|
||||
|
||||
set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wall -std=c++11")
|
||||
include_directories(${CMAKE_CURRENT_SOURCE_DIR}/include ${CUDA_INCLUDE_DIRS} ${OPENCV_INCLUDE_DIRS} ${NVINFER_INCLUDES})
|
||||
|
||||
Reference in New Issue
Block a user