27 lines
1.1 KiB
Bash
Executable file
27 lines
1.1 KiB
Bash
Executable file
#!/bin/sh
|
|
set -e
|
|
|
|
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 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
|