mikegerber/test-nvidia
1
0
Fork 0
Code Issues Pull requests Releases Wiki Activity

⬆️ Update for TensorFlow 2.5+2.6

Browse source
This commit is contained in:
Gerber, Mike 2021-10-20 11:32:40 +02:00
parent 5d05953b5f
commit 4d357d4d0d
2 changed files with 33 additions and 12 deletions
Download patch file Download diff file Expand all files Collapse all files

25
run-docker
View file

@ -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