Skip to content
Advertisement

The most robust way to work with the specific version of Tensorflow with GPU support

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.

User contributions licensed under: CC BY-SA
5 People found this is helpful
Advertisement