Merge pull request #3340 from lissyx/tf-1.15.4-r0.8

Fix #3088: Use TensorFlow 1.15.4 with CUDNN fix
mozilla · Sep 25, 2020 · 3ebcbde · 3ebcbde
2 parents 9ff6df7 + 62c1075
commit 3ebcbde
Show file tree

Hide file tree

Showing 4 changed files with 6 additions and 3 deletions.
diff --git a/Dockerfile.build.tmpl b/Dockerfile.build.tmpl
@@ -64,7 +64,7 @@ ENV TF_CUDA_PATHS "/usr,/usr/local/cuda-10.1,/usr/lib/x86_64-linux-gnu/"
 ENV TF_CUDA_VERSION 10.1
 ENV TF_CUDNN_VERSION 7.6
 ENV TF_CUDA_COMPUTE_CAPABILITIES 6.0
-ENV TF_NCCL_VERSION 2.4
+ENV TF_NCCL_VERSION 2.7
 
 # Common Environment Setup
 ENV TF_BUILD_CONTAINER_TYPE GPU

diff --git a/Dockerfile.train.tmpl b/Dockerfile.train.tmpl
@@ -43,6 +43,9 @@ RUN pip3 install --upgrade native_client/ctcdecode/dist/*.whl
 # Prepare deps
 RUN pip3 install --upgrade pip==20.2.2 wheel==0.34.2 setuptools==49.6.0
 
+# FIXME: Hack because no 1.15.4 Docker image
+RUN pip3 install --upgrade tensorflow-gpu==1.15.4
+
 # Install DeepSpeech
 #  - No need for the decoder since we did it earlier
 #  - There is already correct TensorFlow GPU installed on the base image,

diff --git a/doc/TRAINING.rst b/doc/TRAINING.rst
@@ -65,7 +65,7 @@ If you have a capable (NVIDIA, at least 8GB of VRAM) GPU, it is highly recommend
 .. code-block:: bash
 
    pip3 uninstall tensorflow
-   pip3 install 'tensorflow-gpu==1.15.2'
+   pip3 install 'tensorflow-gpu==1.15.4'
 
 Please ensure you have the required :ref:`CUDA dependency <cuda-deps>`.
 

diff --git a/setup.py b/setup.py
@@ -74,7 +74,7 @@ def main():
     ]
 
     tensorflow_pypi_dep = [
-        'tensorflow == 1.15.2'
+        'tensorflow == 1.15.4'
     ]
 
     # Due to pip craziness environment variables are the only consistent way to