Tensorflow with GPU on Linux or WSL2

After tensorflow 2.10 you can’t use tensorflow-gpu on the Window OS so you need to use WSL on Window 10 or Window 11 to create the conda environment to run tensorflow with your GPU.

Install WSL2

open Window PowerShell and run this command

wsl --install

and then restart PC, open Ubuntu App and do following these next steps.

Install Miniconda

curl https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh -o Miniconda3-latest-Linux-x86_64.sh
bash Miniconda3-latest-Linux-x86_64.sh

press enter until the end and type “yes”

type “yes” and restart your terminal

Create a conda environment

conda create --name tf python=3.9

Activate a conda environment

conda activate tf

Deactivate a conda environment (optional)

conda deactivate

Install NVIDIA GPU driver

https://www.nvidia.com/Download/index.aspx

Verify installed NVIDIA GPU driver (optional)

nvidia-smi

Install CUDA and cuDNN with conda and pip

conda install -c conda-forge cudatoolkit=11.8.0
pip install nvidia-cudnn-cu11==8.6.0.163

You can do it with the following command every time you start a new terminal after activating your conda environment.

CUDNN_PATH=$(dirname $(python -c "import nvidia.cudnn;print(nvidia.cudnn.__file__)"))
export LD_LIBRARY_PATH=$CONDA_PREFIX/lib/:$CUDNN_PATH/lib:$LD_LIBRARY_PATH

The system paths will be automatically configured when you activate this conda environment.

mkdir -p $CONDA_PREFIX/etc/conda/activate.d
echo 'CUDNN_PATH=$(dirname $(python -c "import nvidia.cudnn;print(nvidia.cudnn.__file__)"))' >> $CONDA_PREFIX/etc/conda/activate.d/env_vars.sh
echo 'export LD_LIBRARY_PATH=$CONDA_PREFIX/lib/:$CUDNN_PATH/lib:$LD_LIBRARY_PATH' >> $CONDA_PREFIX/etc/conda/activate.d/env_vars.sh

Install TensorFlow

pip install --upgrade pip
pip install tensorflow==2.12.*

Verify the CPU setup (optional)

python3 -c "import tensorflow as tf; print(tf.reduce_sum(tf.random.normal([1000, 1000])))"

Verify the GPU setup (optional)

python3 -c "import tensorflow as tf; print(tf.config.list_physical_devices('GPU'))"

In Ubuntu 22.04, you may get some error (optional)

Can't find libdevice directory ${CUDA_DIR}/nvvm/libdevice.
...
Couldn't invoke ptxas --version
...
InternalError: libdevice not found at ./libdevice.10.bc [Op:__some_op]

you can fix the error using these commands.

# Install NVCC
conda install -c nvidia cuda-nvcc=11.3.58
# Configure the XLA cuda directory
mkdir -p $CONDA_PREFIX/etc/conda/activate.d
printf 'export XLA_FLAGS=--xla_gpu_cuda_data_dir=$CONDA_PREFIX/lib/\n' >> $CONDA_PREFIX/etc/conda/activate.d/env_vars.sh
source $CONDA_PREFIX/etc/conda/activate.d/env_vars.sh
# Copy libdevice file to the required path
mkdir -p $CONDA_PREFIX/lib/nvvm/libdevice
cp $CONDA_PREFIX/lib/libdevice.10.bc $CONDA_PREFIX/lib/nvvm/libdevice/

Check Ubuntu Version (optional)

lsb_release -a

Directly access Ubuntu resources

References

Install TensorFlow with pip

Install TensorFlow with pip