Manjunath Kudlur | f41959c | 2015-11-06 16:27:58 -0800 | [diff] [blame] | 1 | #!/bin/bash |
| 2 | |
| 3 | ## Set up Cuda-related environment settings |
| 4 | |
| 5 | while [ "$TF_NEED_CUDA" == "" ]; do |
Manjunath Kudlur | 3b8b69f | 2015-11-08 21:41:37 -0800 | [diff] [blame] | 6 | read -p "Do you wish to build TensorFlow with GPU support? [y/n] " INPUT |
Manjunath Kudlur | f41959c | 2015-11-06 16:27:58 -0800 | [diff] [blame] | 7 | case $INPUT in |
| 8 | [Yy]* ) echo -e "GPU support will be enabled for TensorFlow\n"; TF_NEED_CUDA=1;; |
| 9 | [Nn]* ) echo -e "No GPU support will be enabled for TensorFlow\n"; TF_NEED_CUDA=0;; |
| 10 | * ) echo "Invalid selection: " $INPUT;; |
| 11 | esac |
| 12 | done |
| 13 | |
| 14 | if [ "$TF_NEED_CUDA" == "0" ]; then |
| 15 | echo "Configuration finished" |
| 16 | exit |
| 17 | fi |
| 18 | |
| 19 | # Find out where the CUDA toolkit is installed |
| 20 | while true; do |
| 21 | fromuser="" |
| 22 | if [ -z "$CUDA_TOOLKIT_PATH" ]; then |
| 23 | default_cuda_path=/usr/local/cuda |
| 24 | read -p "Please specify the location where CUDA 7.0 toolkit is installed. Refer to README.md for more details. [Default is $default_cuda_path]: " CUDA_TOOLKIT_PATH |
| 25 | fromuser="1" |
| 26 | if [ -z "$CUDA_TOOLKIT_PATH" ]; then |
| 27 | CUDA_TOOLKIT_PATH=$default_cuda_path |
| 28 | fi |
| 29 | fi |
| 30 | if [ -e "$CUDA_TOOLKIT_PATH/lib64/libcudart.so.7.0" ]; then |
| 31 | break |
| 32 | fi |
| 33 | echo "Invalid path to CUDA 7.0 toolkit. ${CUDA_TOOLKIT_PATH}/lib64/libcudart.so.7.0 cannot be found" |
| 34 | if [ -z "$fromuser" ]; then |
| 35 | exit 1 |
| 36 | fi |
| 37 | CUDA_TOOLKIT_PATH="" |
| 38 | # Retry |
| 39 | done |
| 40 | |
| 41 | # Find out where the CUDNN library is installed |
| 42 | while true; do |
| 43 | fromuser="" |
| 44 | if [ -z "$CUDNN_INSTALL_PATH" ]; then |
| 45 | default_cudnn_path=${CUDA_TOOLKIT_PATH} |
| 46 | read -p "Please specify the location where CUDNN 6.5 V2 library is installed. Refer to README.md for more details. [Default is $default_cudnn_path]: " CUDNN_INSTALL_PATH |
| 47 | fromuser="1" |
| 48 | if [ -z "$CUDNN_INSTALL_PATH" ]; then |
| 49 | CUDNN_INSTALL_PATH=$default_cudnn_path |
| 50 | fi |
| 51 | # Result returned from "read" will be used unexpanded. That make "~" unuseable. |
| 52 | # Going through one more level of expansion to handle that. |
| 53 | CUDNN_INSTALL_PATH=$(bash -c "readlink -f $CUDNN_INSTALL_PATH") |
| 54 | fi |
| 55 | if [ -e "$CUDNN_INSTALL_PATH/libcudnn.so.6.5" -o -e "$CUDNN_INSTALL_PATH/lib64/libcudnn.so.6.5" ]; then |
| 56 | break |
| 57 | fi |
| 58 | echo "Invalid path to CUDNN 6.5 V2 toolkit. Neither of the following two files can be found:" |
| 59 | echo "$CUDNN_INSTALL_PATH/lib64/libcudnn.so.6.5" |
| 60 | echo "$CUDNN_INSTALL_PATH/libcudnn.so.6.5" |
| 61 | if [ -z "$fromuser" ]; then |
| 62 | exit 1 |
| 63 | fi |
| 64 | CUDNN_INSTALL_PATH="" |
| 65 | # Retry |
| 66 | done |
| 67 | |
| 68 | cat > third_party/gpus/cuda/cuda.config <<EOF |
| 69 | # CUDA_TOOLKIT_PATH refers to the CUDA toolkit. Tensorflow requries Cuda 7.0 |
| 70 | # at the moment. |
| 71 | CUDA_TOOLKIT_PATH="$CUDA_TOOLKIT_PATH" |
| 72 | |
| 73 | # CUDNN_INSTALL_PATH refers to the CUDNN toolkit. The cudnn header and library |
| 74 | # files can be either in this directory, or under include/ and lib64/ |
| 75 | # directories separately. |
| 76 | CUDNN_INSTALL_PATH="$CUDNN_INSTALL_PATH" |
| 77 | EOF |
| 78 | |
| 79 | # Invoke the cuda_config.sh and set up the TensorFlow's canonical view of the Cuda libraries |
| 80 | (cd third_party/gpus/cuda; ./cuda_config.sh;) || exit -1 |
| 81 | |
| 82 | echo "Configuration finished" |