⬆️ Update for TensorFlow 2.5+2.6
This commit is contained in:
parent
5d05953b5f
commit
4d357d4d0d
2 changed files with 33 additions and 12 deletions
20
README.md
20
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
|
||||
~~~
|
||||
|
||||
|
|
25
run-docker
25
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
|
||||
|
|
Loading…
Add table
Add a link
Reference in a new issue