tensorflow_addons

发表于 更新于 阅读次数：

Build tensorflow/addons with custom cuda/cudnn version

There have four steps to “Build tensorflow/addons with custom cuda/cudnn version”.

your can use conda to duild.

1
2
conda create -n mk
conda activate mk

Requirement

then use conda/pip to install some pkgs.

  • pip
  • six
  • numpy
  • wheel
  • setuptools
  • mock future>=0.17.1
  • keras_applications==1.0.6
  • keras_preprocessing==1.0.5
  • bazel=0.24.1

First

Pull the source code from the github.

1
2
3
4
5
6
7
git clone https://github.com/tensorflow/addons.git
cd addons
# we should checkout to Latest release version  
git checkout r0.6   
# then clean bazel buffer
bazel clean --expunge

Second

Then, in the addons folder，we need to change ‘configure.sh’ and ‘setup.py’.

configure.sh

1
vim configure.sh

we need to alter line 82 to 87 to match your cuda/cudnn version.

1
2
3
4
5
6
7
8

if [[ "$TF_NEED_CUDA" == "1" ]]; then
    write_action_env_to_bazelrc "TF_NEED_CUDA" ${TF_NEED_CUDA}
    write_action_env_to_bazelrc "CUDNN_INSTALL_PATH" "${CUDNN_INSTALL_PATH:=/disk1/lx/cuda/lib64}
    write_action_env_to_bazelrc "TF_CUDA_VERSION" "9.0"
    write_action_env_to_bazelrc "TF_CUDNN_VERSION" "7"
    write_action_env_to_bazelrc "CUDA_TOOLKIT_PATH" "${CUDA_HOME:=/disk1/lx/cuda}"

setup.py

1
vim setup.py

We can specify the TF2 version you need in line 69~80

1
2
3
4
5
6
7
8
                                                                                 
if project_name == TFA_RELEASE:                                                  
    # TODO: remove if-else condition when tf supports package consolidation.     
    if platform.system() == 'Linux':                                             
        REQUIRED_PACKAGES.append('tensorflow == 2.0.0')                          
    else:                                                                        
        REQUIRED_PACKAGES.append('tensorflow == 2.0.0')

Third

change addons/build_deps/ requirements_gpu.txt and requirements.txt to match your TF2 version.

requirements_gpu.txt

1
2
# TensorFlow greater than this date is manylinux2010 compliant                             
# nothing

requirements.txt

1
2
# TensorFlow greater than this date is manylinux2010 compliant             
tensorflow==2.0.0

Forth

bazel and install

1
2
3
4
5
6
7
8

# This script links project with TensorFlow dependency
./configure.sh

bazel build build_pip_pkg
bazel-bin/build_pip_pkg artifacts 

pip install artifacts/tensorflow_addons-*.whl