TensorRT version

This commit is contained in:
Francesco Gatti
2017-08-01 17:51:49 +02:00
parent 1cfe70365f
commit ed5e5d58b5
6 changed files with 16 additions and 11 deletions
+2 -1
View File
@@ -11,11 +11,12 @@ cuda_add_library(kernels SHARED src/kernels/activation_elu.cu
src/kernels/reorg.cu
src/kernels/softmax.cu)
set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wall -std=c++11")
include_directories(${CMAKE_CURRENT_SOURCE_DIR}/include ${CUDA_INCLUDE_DIRS})
add_library(tkDNN SHARED src/Layer.cpp src/LayerWgs.cpp
src/Dense.cpp src/Activation.cpp src/Conv2d.cpp src/Flatten.cpp src/MulAdd.cpp src/Pooling.cpp src/Softmax.cpp
src/Route.cpp src/Reorg.cpp src/Region.cpp src/Network.cpp src/utils.cpp)
target_link_libraries(tkDNN kernels ${CUDA_LIBRARIES} ${CUDA_CUBLAS_LIBRARIES} -lcudnn)
target_link_libraries(tkDNN kernels ${CUDA_LIBRARIES} ${CUDA_CUBLAS_LIBRARIES} -lcudnn -lnvinfer)
add_executable(test_simple tests/test/test.cpp)
target_link_libraries(test_simple tkDNN)