Upgrade to Tensorflow 2.0 (#124)

b/159364787

Dockerfile

@@ -58,7 +58,7 @@ RUN apt-get install -y libzmq3-dev python-pip default-jdk && \
 # Keras sets up a virtualenv and installs tensorflow
 # in the WORKON_HOME directory, so choose an explicit location for it.
 ENV WORKON_HOME=/usr/local/share/.virtualenvs
-RUN pip install --user virtualenv && R -e 'keras::install_keras(tensorflow = "1.15", extra_packages = c("pandas", "numpy", "pycryptodome"))'
+RUN pip install --user virtualenv && R -e 'keras::install_keras(tensorflow = "2.0", extra_packages = c("pandas", "numpy", "pycryptodome"))'
 
 # Install kaggle libraries.
 # Do this at the end to avoid rebuilding everything when any change is made.

gpu.Dockerfile

@@ -1,5 +1,5 @@
 ARG BASE_TAG=staging
-FROM nvidia/cuda:10.0-cudnn7-devel-ubuntu16.04 AS nvidia
+FROM nvidia/cuda:10.0-cudnn7-devel-ubuntu18.04 AS nvidia
 FROM gcr.io/kaggle-images/rstats:${BASE_TAG}
 ARG ncpus=1
 
@@ -38,7 +38,9 @@ RUN apt-get update && apt-get install -y --no-install-recommends \
       cuda-minimal-build-$CUDA_PKG_VERSION \
       cuda-command-line-tools-$CUDA_PKG_VERSION \
       libcudnn7=$CUDNN_VERSION-1+cuda$CUDA_MAJOR_VERSION.$CUDA_MINOR_VERSION \
-      libcudnn7-dev=$CUDNN_VERSION-1+cuda$CUDA_MAJOR_VERSION.$CUDA_MINOR_VERSION && \
+      libcudnn7-dev=$CUDNN_VERSION-1+cuda$CUDA_MAJOR_VERSION.$CUDA_MINOR_VERSION  \
+      libnccl2=2.5.6-1+cuda$CUDA_MAJOR_VERSION.$CUDA_MINOR_VERSION \
+      libnccl-dev=2.5.6-1+cuda$CUDA_MAJOR_VERSION.$CUDA_MINOR_VERSION && \
     ln -s /usr/local/cuda-$CUDA_MAJOR_VERSION.$CUDA_MINOR_VERSION /usr/local/cuda && \
     ln -s /usr/local/cuda/lib64/stubs/libcuda.so /usr/local/cuda/lib64/stubs/libcuda.so.1 && \
     /tmp/clean-layer.sh
@@ -53,7 +55,7 @@ ENV CUDA_HOME=/usr/local/cuda
 ADD ldpaths $R_HOME/etc/ldpaths
 
 # Install tensorflow with GPU support
-RUN R -e 'keras::install_keras(tensorflow = "1.15-gpu")' && \
+RUN R -e 'keras::install_keras(tensorflow = "2.0-gpu")' && \
     rm -rf /tmp/tensorflow_gpu && \
     /tmp/clean-layer.sh
 

tests/test_tensorflow.R

@@ -2,14 +2,11 @@ context("tensorflow")
 
 test_that("gpu imports", {
   check_gpu()
-
-  #TODO(b/143642025): update to TF2.0 compatible and fix potential CUDA issues.
 
   library(tensorflow)
-  sess <- tf$Session()
   with(tf$device("/gpu:0"), {
       const <- tf$constant(42)
-      expect_equal(42, sess$run(const))
+      expect_equal(42, as.integer(const))
   })
 
   expect_true(TRUE)