Local GPU

The default build of TensorFlow will use an NVIDIA® GPU if it is available and the appropriate drivers are installed, and otherwise fallback to using the CPU only.

The prerequisites for the GPU version of TensorFlow on each platform are covered below. Note that on all platforms (except macOS) you must be running an NVIDIA® GPU with CUDA® Compute Capability 3.5 or higher. See the list of CUDA-enabled GPU cards.

To enable TensorFlow to use a local NVIDIA® GPU, you can install the following:

To install the required NVIDIA components on Ubuntu 20.04, you can run the following at the terminal:

wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2004/x86_64/cuda-ubuntu2004.pin
sudo mv cuda-ubuntu2004.pin /etc/apt/preferences.d/cuda-repository-pin-600
sudo apt-key adv --fetch-keys https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2004/x86_64/7fa2af80.pub
sudo add-apt-repository "deb https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2004/x86_64/ /"
sudo apt-get update
sudo apt-get -y install cuda-11-2 libcudnn8=8.1.1.33-1+cuda11.2
sudo apt-mark hold libcudnn8

First, confirm you have an appropriate GPU. This article describes how to detect whether your graphics card uses an NVIDIA® GPU:

http://nvidia.custhelp.com/app/answers/detail/a_id/2040/~/identifying-the-graphics-card-model-and-device-id-in-a-pc

Once you’ve confirmed that you have an NVIDIA® GPU, the following article describes how to install required software components including the CUDA Toolkit v11.2 and cuDNN 8.1:

https://www.tensorflow.org/install/gpu#hardware_requirements

Note that the documentation on installation of the last component (cuDNN) is a bit sparse. Once you join the NVIDIA® developer program and download the zip file containing cuDNN you need to extract the zip file and add the location where you extracted it to your system PATH.

The specifics of installing required software differ by Linux version so please review the NVIDIA® documentation carefully to ensure you install everything correctly.

TensorFlow does not support usage of NVIDIA GPUs on Macs. For ARM (M1, M2) Macs however, Apple provides a custom build of TensorFlow that can take advantage of the GPU. install_tensorflow() will detect if it is being run on an Arm Mac and automatically install the appropriate Python packages (tensorflow-macos and tensorflow-metal).

Conda is the only supported installation method on Arm Macs. Additionally, x86_64 builds of R running under Rosetta are not supported; an Arm64 build of R is required.

You can test if you have an Arm build of R like this:

R.version$arch

If this returns "aarch64", you’re using the correct build of R. If instead it returns "x86_64", then you need to install a different build of R. You can install an Arm build of R by navigating to https://cloud.r-project.org/bin/macosx/ and selecting the appropriate package (e.g., “R-4.2.1-arm64.pkg”) or using an installation manager like rig and running rig add release.

You can configure reticulate to use a conda Python using reticulate like this:

reticulate::install_miniconda()
tensorflow::install_tensorflow()

Note that install_miniconda() automatically creates a conda environment named r-reticulate, that becomes the default environment used by reticulate.

Test

You can test if TensorFlow can use a local GPU by running the following:

library(tensorflow)
tf$config$list_physical_devices("GPU")

If you see a list like this:

[[1]]
PhysicalDevice(name='/physical_device:GPU:0', device_type='GPU')

Then that means TensorFlow is configured properly. If instead you see an empty list(), then TensorFlow is not using a GPU.

For support with installation, open a Github issue