The project I am working on uses the Tensorflow 2.0
, which is not the most recent version of this Deep Learning Framework and I would like to enable GPU support.
I have already locally installed Tensorflow 2.4.1
for another purposes with CUDA-11.0
. This version of Tensorflow
sees the GPU on my PC and I can perform training and inference without problem. However, another version, 2.0
, installed in the virtualenv
doesn’t see the GPU.
As far as I understood, it is due to the fact, that 2.0
works only with CUDA-10.0
.
What would be the most robust way to setup Tensorflow 2.0 + CUDA-10.0
. Ordinary installation of CUDA Toolkit
seems like it would overwrite the existing 11-0
installation, which is undesirable.
I’m not familiar a lot with the containers. Would it be a good solution to create a Docker
image with that specific version of Tensorflow
and CUDA
?
I would be grateful for suggestions and comments.
Advertisement
Answer
Definitely explore the tensorflow docker images: https://hub.docker.com/r/tensorflow/tensorflow/tags
There are a good range of precompiled images which work with gpu, providing you set --runtime=nvidia
parameter on docker run
Some resources on the nvidia runtime: https://www.tensorflow.org/install/docker https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/user-guide.html
the nvidia docs above are bit complicated looking. On Debian I was able to install with apt-get install nvidia-container-runtime
(more info: https://github.com/NVIDIA/nvidia-container-runtime). Setting the nvidia runtime as default is not required.
You can then launch a GPU enabled tensorflow container with:
docker run -it --rm --runtime=nvidia -v $(pwd):/my-src -u $(id -u):$(id -g) tensorflow/tensorflow:2.2.2-gpu /bin/bash
inside the container:
python -c "import tensorflow"
will verify the setup.
In practice, I’ve always been able to use the latest nvidia drivers even with older tensorflow versions. At the moment, I’m flipping between tf 2.2.2 and tf 2.4.1 with nvidia 460 drivers and CUDA 11.2. I only had version issues when I was trying to use tf 1.4, which is pretty old now.