Skip to content

Use official TensorFlow wheel #1036

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 3 commits into from
Jun 25, 2021
Merged

Use official TensorFlow wheel #1036

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
91eae35
Install official tensorflow whl.
rosbo Jun 23, 2021
8fdf1bb
Reclare TENSORFLOW_VERSION ARG to be used in commands after FROM
rosbo Jun 23, 2021
50c5305
Install proper version for tensorflow_gcs_config
rosbo Jun 25, 2021
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
27 changes: 7 additions & 20 deletions Dockerfile
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,9 +1,12 @@
ARG BASE_TAG=m73
ARG TENSORFLOW_VERSION=2.4.1

FROM gcr.io/kaggle-images/python-tensorflow-whl:${TENSORFLOW_VERSION}-py37-2 as tensorflow_whl
FROM gcr.io/deeplearning-platform-release/base-cpu:${BASE_TAG}

# We need to redefine TENSORFLOW_VERSION here to get the default ARG value defined above the FROM instruction.
# See: https://docs.docker.com/engine/reference/builder/#understand-how-arg-and-from-interact
ARG TENSORFLOW_VERSION

ADD clean-layer.sh /tmp/clean-layer.sh
ADD patches/nbconvert-extensions.tpl /opt/kaggle/nbconvert-extensions.tpl
ADD patches/template_conf.json /opt/kaggle/conf.json
Expand Down Expand Up @@ -56,22 +59,9 @@ RUN pip install seaborn python-dateutil dask && \
pip install -f https://h2o-release.s3.amazonaws.com/h2o/latest_stable_Py.html h2o && \
/tmp/clean-layer.sh

# Install tensorflow from a pre-built wheel
COPY --from=tensorflow_whl /tmp/tensorflow_cpu/*.whl /tmp/tensorflow_cpu/
RUN pip install /tmp/tensorflow_cpu/tensorflow*.whl && \
rm -rf /tmp/tensorflow_cpu && \
/tmp/clean-layer.sh

# Install tensorflow-gcs-config from a pre-built wheel
COPY --from=tensorflow_whl /tmp/tensorflow_gcs_config/*.whl /tmp/tensorflow_gcs_config/
RUN pip install /tmp/tensorflow_gcs_config/tensorflow*.whl && \
rm -rf /tmp/tensorflow_gcs_config && \
/tmp/clean-layer.sh

# Install TensorFlow addons (TFA).
COPY --from=tensorflow_whl /tmp/tfa_cpu/*.whl /tmp/tfa_cpu/
RUN pip install /tmp/tfa_cpu/tensorflow*.whl && \
rm -rf /tmp/tfa_cpu/ && \
RUN pip install tensorflow==${TENSORFLOW_VERSION} && \
pip install tensorflow-gcs-config==2.4.0 && \
pip install tensorflow-addons==0.12.1 && \
/tmp/clean-layer.sh

RUN apt-get install -y libfreetype6-dev && \
Expand Down Expand Up @@ -513,9 +503,6 @@ RUN jupyter-nbextension disable nb_conda --py --sys-prefix && \
# Set backend for matplotlib
ENV MPLBACKEND "agg"

# We need to redefine TENSORFLOW_VERSION here to get the default ARG value defined above the FROM instruction.
# See: https://docs.docker.com/engine/reference/builder/#understand-how-arg-and-from-interact
ARG TENSORFLOW_VERSION
ARG GIT_COMMIT=unknown
ARG BUILD_DATE=unknown

Expand Down
11 changes: 1 addition & 10 deletions gpu.Dockerfile
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,7 +1,6 @@
ARG BASE_TAG=staging

FROM nvidia/cuda:11.0-cudnn8-devel-ubuntu18.04 AS nvidia
FROM gcr.io/kaggle-images/python-tensorflow-whl:2.4.1-py37-2 as tensorflow_whl
FROM gcr.io/kaggle-images/python:${BASE_TAG}

ADD clean-layer.sh /tmp/clean-layer.sh
Expand Down Expand Up @@ -82,20 +81,12 @@ RUN pip install jax==0.2.12 jaxlib==0.1.64+cuda$CUDA_MAJOR_VERSION$CUDA_MINOR_V
/tmp/clean-layer.sh

# Reinstall packages with a separate version for GPU support.
COPY --from=tensorflow_whl /tmp/tensorflow_gpu/*.whl /tmp/tensorflow_gpu/
RUN pip uninstall -y tensorflow && \
pip install /tmp/tensorflow_gpu/tensorflow*.whl && \
rm -rf /tmp/tensorflow_gpu && \
pip install tensorflow-gpu==2.4.1 && \
pip uninstall -y mxnet && \
pip install mxnet-cu$CUDA_MAJOR_VERSION$CUDA_MINOR_VERSION && \
/tmp/clean-layer.sh

# Reinstall TensorFlow addons (TFA) with GPU support.
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

No need to reinstall TFA, the official whl supports CPU & GPU.

COPY --from=tensorflow_whl /tmp/tfa_gpu/*.whl /tmp/tfa_gpu/
RUN pip install /tmp/tfa_gpu/tensorflow*.whl && \
rm -rf /tmp/tfa_gpu/ && \
/tmp/clean-layer.sh

# Install GPU-only packages
RUN pip install pycuda && \
pip install cupy-cuda$CUDA_MAJOR_VERSION$CUDA_MINOR_VERSION && \
Expand Down