Commit 3d698a6
authored
Use official TensorFlow wheel (Kaggle#1036)
We are moving towards the base DLVM TF image.
http://b/1866460691 parent 0b2dcfa commit 3d698a6
2 files changed
Lines changed: 8 additions & 30 deletions
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
@@ -1,9 +1,12 @@ | |||
| 1 | 1 | ARG BASE_TAG=m73 | |
| 2 | 2 | ARG TENSORFLOW_VERSION=2.4.1 | |
| 3 | 3 | ||
| 4 | - FROM gcr.io/kaggle-images/python-tensorflow-whl:${TENSORFLOW_VERSION}-py37-2 as tensorflow_whl | ||
| 5 | 4 | FROM gcr.io/deeplearning-platform-release/base-cpu:${BASE_TAG} | |
| 6 | 5 | ||
| 6 | + # We need to redefine TENSORFLOW_VERSION here to get the default ARG value defined above the FROM instruction. | ||
| 7 | + # See: https://docs.docker.com/engine/reference/builder/#understand-how-arg-and-from-interact | ||
| 8 | + ARG TENSORFLOW_VERSION | ||
| 9 | + | ||
| 7 | 10 | ADD clean-layer.sh /tmp/clean-layer.sh | |
| 8 | 11 | ADD patches/nbconvert-extensions.tpl /opt/kaggle/nbconvert-extensions.tpl | |
| 9 | 12 | ADD patches/template_conf.json /opt/kaggle/conf.json | |
@@ -56,22 +59,9 @@ RUN pip install seaborn python-dateutil dask && \ | |||
| 56 | 59 | pip install -f https://h2o-release.s3.amazonaws.com/h2o/latest_stable_Py.html h2o && \ | |
| 57 | 60 | /tmp/clean-layer.sh | |
| 58 | 61 | ||
| 59 | - # Install tensorflow from a pre-built wheel | ||
| 60 | - COPY --from=tensorflow_whl /tmp/tensorflow_cpu/*.whl /tmp/tensorflow_cpu/ | ||
| 61 | - RUN pip install /tmp/tensorflow_cpu/tensorflow*.whl && \ | ||
| 62 | - rm -rf /tmp/tensorflow_cpu && \ | ||
| 63 | - /tmp/clean-layer.sh | ||
| 64 | - | ||
| 65 | - # Install tensorflow-gcs-config from a pre-built wheel | ||
| 66 | - COPY --from=tensorflow_whl /tmp/tensorflow_gcs_config/*.whl /tmp/tensorflow_gcs_config/ | ||
| 67 | - RUN pip install /tmp/tensorflow_gcs_config/tensorflow*.whl && \ | ||
| 68 | - rm -rf /tmp/tensorflow_gcs_config && \ | ||
| 69 | - /tmp/clean-layer.sh | ||
| 70 | - | ||
| 71 | - # Install TensorFlow addons (TFA). | ||
| 72 | - COPY --from=tensorflow_whl /tmp/tfa_cpu/*.whl /tmp/tfa_cpu/ | ||
| 73 | - RUN pip install /tmp/tfa_cpu/tensorflow*.whl && \ | ||
| 74 | - rm -rf /tmp/tfa_cpu/ && \ | ||
| 62 | + RUN pip install tensorflow==${TENSORFLOW_VERSION} && \ | ||
| 63 | + pip install tensorflow-gcs-config==2.4.0 && \ | ||
| 64 | + pip install tensorflow-addons==0.12.1 && \ | ||
| 75 | 65 | /tmp/clean-layer.sh | |
| 76 | 66 | ||
| 77 | 67 | RUN apt-get install -y libfreetype6-dev && \ | |
@@ -513,9 +503,6 @@ RUN jupyter-nbextension disable nb_conda --py --sys-prefix && \ | |||
| 513 | 503 | # Set backend for matplotlib | |
| 514 | 504 | ENV MPLBACKEND "agg" | |
| 515 | 505 | ||
| 516 | - # We need to redefine TENSORFLOW_VERSION here to get the default ARG value defined above the FROM instruction. | ||
| 517 | - # See: https://docs.docker.com/engine/reference/builder/#understand-how-arg-and-from-interact | ||
| 518 | - ARG TENSORFLOW_VERSION | ||
| 519 | 506 | ARG GIT_COMMIT=unknown | |
| 520 | 507 | ARG BUILD_DATE=unknown | |
| 521 | 508 | ||
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
@@ -1,7 +1,6 @@ | |||
| 1 | 1 | ARG BASE_TAG=staging | |
| 2 | 2 | ||
| 3 | 3 | FROM nvidia/cuda:11.0-cudnn8-devel-ubuntu18.04 AS nvidia | |
| 4 | - FROM gcr.io/kaggle-images/python-tensorflow-whl:2.4.1-py37-2 as tensorflow_whl | ||
| 5 | 4 | FROM gcr.io/kaggle-images/python:${BASE_TAG} | |
| 6 | 5 | ||
| 7 | 6 | ADD clean-layer.sh /tmp/clean-layer.sh | |
@@ -82,20 +81,12 @@ RUN pip install jax==0.2.12 jaxlib==0.1.64+cuda$CUDA_MAJOR_VERSION$CUDA_MINOR_V | |||
| 82 | 81 | /tmp/clean-layer.sh | |
| 83 | 82 | ||
| 84 | 83 | # Reinstall packages with a separate version for GPU support. | |
| 85 | - COPY --from=tensorflow_whl /tmp/tensorflow_gpu/*.whl /tmp/tensorflow_gpu/ | ||
| 86 | 84 | RUN pip uninstall -y tensorflow && \ | |
| 87 | - pip install /tmp/tensorflow_gpu/tensorflow*.whl && \ | ||
| 88 | - rm -rf /tmp/tensorflow_gpu && \ | ||
| 85 | + pip install tensorflow-gpu==2.4.1 && \ | ||
| 89 | 86 | pip uninstall -y mxnet && \ | |
| 90 | 87 | pip install mxnet-cu$CUDA_MAJOR_VERSION$CUDA_MINOR_VERSION && \ | |
| 91 | 88 | /tmp/clean-layer.sh | |
| 92 | 89 | ||
| 93 | - # Reinstall TensorFlow addons (TFA) with GPU support. | ||
| 94 | - COPY --from=tensorflow_whl /tmp/tfa_gpu/*.whl /tmp/tfa_gpu/ | ||
| 95 | - RUN pip install /tmp/tfa_gpu/tensorflow*.whl && \ | ||
| 96 | - rm -rf /tmp/tfa_gpu/ && \ | ||
| 97 | - /tmp/clean-layer.sh | ||
| 98 | - | ||
| 99 | 90 | # Install GPU-only packages | |
| 100 | 91 | RUN pip install pycuda && \ | |
| 101 | 92 | pip install cupy-cuda$CUDA_MAJOR_VERSION$CUDA_MINOR_VERSION && \ | |
0 commit comments