Fixed and re-enables TensorRT steps (#14960)

* Pins libnvinfer versions * Sets cudnn to version 7.5.0 in tensorrt environment * Re-enables TensorRT stages
apache · May 15, 2019 · f7b7163 · f7b7163 · KellenSunderland · May 15, 2019
1 parent 99f5f66
commit f7b7163
Show file tree

Hide file tree

Showing 4 changed files with 10 additions and 5 deletions.
diff --git a/ci/docker/Dockerfile.build.ubuntu_gpu_tensorrt b/ci/docker/Dockerfile.build.ubuntu_gpu_tensorrt
@@ -18,7 +18,7 @@
 #
 # Dockerfile to run MXNet on Ubuntu 16.04 for CPU
 
-FROM nvidia/cuda:10.0-cudnn7-devel
+FROM nvidia/cuda:10.0-devel
 
 WORKDIR /work/deps
 
@@ -35,6 +35,10 @@ ARG USER_ID=0
 COPY install/ubuntu_adduser.sh /work/
 RUN /work/ubuntu_adduser.sh
 
+ENV CUDNN_VERSION=7.5.0.56
+COPY install/ubuntu_cudnn.sh /work/
+RUN /work/ubuntu_cudnn.sh
+
 COPY runtime_functions.sh /work/
 
 WORKDIR /work/mxnet

diff --git a/ci/docker/install/tensorrt.sh b/ci/docker/install/tensorrt.sh
@@ -44,5 +44,6 @@ echo "TensorRT build enabled. Installing TensorRT."
 wget -qO tensorrt.deb https://developer.download.nvidia.com/compute/machine-learning/repos/ubuntu1604/x86_64/nvinfer-runtime-trt-repo-ubuntu1604-5.0.2-ga-cuda10.0_1-1_amd64.deb
 dpkg -i tensorrt.deb
 apt-get update
-apt-get install -y --allow-downgrades libnvinfer-dev
+apt-get install -y --allow-downgrades libnvinfer5=5.1.2-1+cuda10.0
+apt-get install -y --allow-downgrades libnvinfer-dev=5.1.2-1+cuda10.0
 rm tensorrt.deb
diff --git a/ci/docker/install/ubuntu_cudnn.sh b/ci/docker/install/ubuntu_cudnn.sh
@@ -50,5 +50,5 @@ case ${CUDA_VERSION} in
 		;;
 esac
 
-apt-get install -y libcudnn7=${libcudnn7_version} libcudnn7-dev=${libcudnn7_dev_version}
+apt-get install -y --allow-downgrades libcudnn7=${libcudnn7_version} libcudnn7-dev=${libcudnn7_dev_version}
 
diff --git a/ci/jenkins/Jenkinsfile_unix_gpu b/ci/jenkins/Jenkinsfile_unix_gpu
@@ -39,7 +39,7 @@ core_logic: {
     custom_steps.compile_unix_full_gpu(),
     custom_steps.compile_unix_cmake_mkldnn_gpu(),
     custom_steps.compile_unix_cmake_gpu(),
-    // custom_steps.compile_unix_tensorrt_gpu(),
+    custom_steps.compile_unix_tensorrt_gpu(),
     custom_steps.compile_unix_int64_gpu()
   ]) 
 
@@ -51,7 +51,7 @@ core_logic: {
     custom_steps.test_unix_python2_mkldnn_gpu(),
     custom_steps.test_unix_python3_mkldnn_gpu(),
     custom_steps.test_unix_python3_mkldnn_nocudnn_gpu(),
-    // custom_steps.test_unix_python3_tensorrt_gpu(),
+    custom_steps.test_unix_python3_tensorrt_gpu(),
     custom_steps.test_unix_perl_gpu(),
     custom_steps.test_unix_r_gpu(),
     custom_steps.test_unix_cpp_gpu(),