From 4d357d4d0dc91e10f926741c24121d33d9c172a6 Mon Sep 17 00:00:00 2001
From: "Gerber, Mike" <mike.gerber@sbb.spk-berlin.de>
Date: Wed, 20 Oct 2021 11:32:40 +0200
Subject: [PATCH] =?UTF-8?q?=E2=AC=86=EF=B8=8F=20Update=20for=20TensorFlow?=
 =?UTF-8?q?=202.5+2.6?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 README.md  | 20 ++++++++++++++------
 run-docker | 25 +++++++++++++++++++------
 2 files changed, 33 insertions(+), 12 deletions(-)

diff --git a/README.md b/README.md
index f66697a..0d84805 100644
--- a/README.md
+++ b/README.md
@@ -13,13 +13,21 @@ Example output
 ==============
 ~~~
 % ./run-docker
-== tf1
-GPU 0: GeForce RTX 2080 (UUID: GPU-612ce75c-1340-772b-039c-2a83a3ea5c95)
-TensorFlow 1.15.3
+== tensorflow-gpu 1.15 nvidia/cuda:10.0-cudnn7-runtime-ubuntu18.04
+GPU 0: NVIDIA GeForce RTX 2080 (UUID: GPU-612ce75c-1340-772b-039c-2a83a3ea5c95)
+TensorFlow 1.15.0
 GPU available: True
-== tf2
-GPU 0: GeForce RTX 2080 (UUID: GPU-612ce75c-1340-772b-039c-2a83a3ea5c95)
-TensorFlow 2.3.0
+== tensorflow 2.4.3 nvidia/cuda:11.0-cudnn8-runtime-ubuntu18.04
+GPU 0: NVIDIA GeForce RTX 2080 (UUID: GPU-612ce75c-1340-772b-039c-2a83a3ea5c95)
+TensorFlow 2.4.3
+GPU available: True
+== tensorflow 2.5.1 nvidia/cuda:11.2.1-cudnn8-runtime-ubuntu18.04
+GPU 0: NVIDIA GeForce RTX 2080 (UUID: GPU-612ce75c-1340-772b-039c-2a83a3ea5c95)
+TensorFlow 2.5.1
+GPU available: True
+== tensorflow 2.6.0 nvidia/cuda:11.2.1-cudnn8-runtime-ubuntu18.04
+GPU 0: NVIDIA GeForce RTX 2080 (UUID: GPU-612ce75c-1340-772b-039c-2a83a3ea5c95)
+TensorFlow 2.6.0
 GPU available: True
 ~~~
 
diff --git a/run-docker b/run-docker
index 584b4c3..4fc05bd 100755
--- a/run-docker
+++ b/run-docker
@@ -1,14 +1,27 @@
 #!/bin/sh
 set -e
 
-for tf in tf1 tf2; do
-  echo "== $tf"
-  case "$tf" in
-    tf1) BASE_IMAGE=nvidia/cuda:10.0-cudnn7-runtime-ubuntu18.04;;
-    tf2) BASE_IMAGE=nvidia/cuda:11.0-cudnn8-runtime-ubuntu18.04;;
+for tensorflow_version in 1.15 2.4.3 2.5.1 2.6.0; do
+
+  BASE_IMAGE=invalid
+  case "$tensorflow_version" in
+    1.15) BASE_IMAGE=nvidia/cuda:10.0-cudnn7-runtime-ubuntu18.04;;
+    2.4.3) BASE_IMAGE=nvidia/cuda:11.0-cudnn8-runtime-ubuntu18.04;;
+    2.5.1) BASE_IMAGE=nvidia/cuda:11.2.1-cudnn8-runtime-ubuntu18.04;;
+    2.6.0) BASE_IMAGE=nvidia/cuda:11.2.1-cudnn8-runtime-ubuntu18.04;;
   esac
+
+  # According to the docs at https://www.tensorflow.org/install/gpu, we should
+  # use different package names depending on the major version of TF.
+  if echo $tensorflow_version | grep -q ^1; then
+    tensorflow_package=tensorflow-gpu
+  else
+     tensorflow_package=tensorflow
+  fi
+
+  echo "== $tensorflow_package $tensorflow_version $BASE_IMAGE"
   work_dir=`dirname $0`
-  image_id=`docker build -q --build-arg tf=$tf --build-arg BASE_IMAGE=$BASE_IMAGE -f assets/Dockerfile $work_dir`
+  image_id=`docker build -q --build-arg tensorflow_package=$tensorflow_package --build-arg tensorflow_version=$tensorflow_version --build-arg BASE_IMAGE=$BASE_IMAGE -f assets/Dockerfile $work_dir`
   docker run --gpus all -it --rm -e TF_CPP_MIN_LOG_LEVEL=1 $image_id
   docker rmi $image_id >/dev/null || true
 done