Compatibilidad con GPU

La compatibilidad con GPU de TensorFlow requiere una selección de controladores y bibliotecas. Para simplificar la instalación y evitar conflictos de bibliotecas, recomendamos usar una imagen de Docker de TensorFlow compatible con GPU (solo Linux). Para esta configuración solo se necesitan los controladores de GPU de NVIDIA®.

Estas instrucciones de instalación corresponden a la actualización más reciente de TensorFlow. Si deseas usar versiones anteriores de TensorFlow, consulta las configuraciones de compilación probadas para las versiones de CUDA® y cuDNN.

Paquete pip

Consulta la guía de instalación de pip para ver los paquetes disponibles, los requisitos del sistema y las instrucciones. El paquete pip de TensorFlow incluye compatibilidad con GPU para tarjetas habilitadas para CUDA®:

pip install tensorflow

En esta guía, se abordan los pasos de instalación y asistencia de GPU para la versión estable más reciente de TensorFlow.

Versiones anteriores de TensorFlow

Para TensorFlow 1.15 y versiones anteriores, los paquetes de CPU y GPU son independientes:

pip install tensorflow==1.15      # CPU
pip install tensorflow-gpu==1.15  # GPU

Requisitos de hardware

Se admiten los siguientes dispositivos habilitados para GPU:

Tarjeta GPU NVIDIA® con arquitecturas CUDA® 3.5, 5.0, 6.0, 7.0, 7.5, 8.0 y versiones posteriores. Consulta la lista de tarjetas GPU habilitadas para CUDA®.
Para usar GPU con arquitecturas CUDA® no compatibles, evitar la compilación de JIT de PTX o usar diferentes versiones de las bibliotecas NVIDIA®, consulta la guía compilación con el código fuente en Linux.
Los paquetes no contienen código PTX, excepto la última arquitectura CUDA® compatible. Por lo tanto, no se puede cargar TensorFlow en las GPU más antiguas cuando se configura CUDA_FORCE_PTX_JIT=1. (Consulta Compatibilidad de aplicaciones para obtener más detalles).

Requisitos de software

Debes instalar el siguiente software de NVIDIA® en tu sistema:

Controladores de GPU NVIDIA®: CUDA® 11.2 requiere la versión 450.80.02 o una posterior.
Kit de herramientas CUDA®: TensorFlow es compatible con CUDA® 11.2 (TensorFlow >= 2.5.0).
CUPTI incluye el kit de herramientas CUDA®.
SDK de cuDNN 8.1.0 (versiones de cuDNN).
(Opcional) TensorRT 6.0 para mejorar la latencia y la capacidad de procesamiento para la inferencia en algunos modelos.

Configuración en Linux

Las instrucciones apt presentadas a continuación constituyen la forma más fácil de instalar el software NVIDIA requerido en Ubuntu. Sin embargo, si compilas TensorFlow con el código fuente, instala de forma manual los requisitos de software mencionados más arriba y procura usar una imagen de Docker de TensorFlow -devel como base.

Instala CUPTI, que incluye el kit de herramientas CUDA®. Adjunta el directorio de instalación de CUPTI a la variable de entorno $LD_LIBRARY_PATH:

export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/cuda/extras/CUPTI/lib64

Instala CUDA con apt

En esta sección se muestra cómo instalar CUDA® 11 (TensorFlow >= 2.4.0) en Ubuntu 16.04 y 18.04. Estas instrucciones pueden funcionar para otras distribuciones basadas en Debian.

Ubuntu 18.04 (CUDA 11.0)

# Add NVIDIA package repositories
wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu1804/x86_64/cuda-ubuntu1804.pin
sudo mv cuda-ubuntu1804.pin /etc/apt/preferences.d/cuda-repository-pin-600
sudo apt-key adv --fetch-keys https://developer.download.nvidia.com/compute/cuda/repos/ubuntu1804/x86_64/7fa2af80.pub
sudo add-apt-repository "deb https://developer.download.nvidia.com/compute/cuda/repos/ubuntu1804/x86_64/ /"
sudo apt-get update

wget http://developer.download.nvidia.com/compute/machine-learning/repos/ubuntu1804/x86_64/nvidia-machine-learning-repo-ubuntu1804_1.0.0-1_amd64.deb

sudo apt install ./nvidia-machine-learning-repo-ubuntu1804_1.0.0-1_amd64.deb
sudo apt-get update

wget https://developer.download.nvidia.com/compute/machine-learning/repos/ubuntu1804/x86_64/libnvinfer7_7.1.3-1+cuda11.0_amd64.deb
sudo apt install ./libnvinfer7_7.1.3-1+cuda11.0_amd64.deb
sudo apt-get update

# Install development and runtime libraries (~4GB)
sudo apt-get install --no-install-recommends \
    cuda-11-0 \
    libcudnn8=8.0.4.30-1+cuda11.0  \
    libcudnn8-dev=8.0.4.30-1+cuda11.0

# Reboot. Check that GPUs are visible using the command: nvidia-smi

# Install TensorRT. Requires that libcudnn8 is installed above.
sudo apt-get install -y --no-install-recommends libnvinfer7=7.1.3-1+cuda11.0 \
    libnvinfer-dev=7.1.3-1+cuda11.0 \
    libnvinfer-plugin7=7.1.3-1+cuda11.0

Ubuntu 16.04 (CUDA 11.0)

# Add NVIDIA package repositories
# Add HTTPS support for apt-key
sudo apt-get install gnupg-curl
wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu1604/x86_64/cuda-ubuntu1604.pin
sudo mv cuda-ubuntu1604.pin /etc/apt/preferences.d/cuda-repository-pin-600
sudo apt-key adv --fetch-keys https://developer.download.nvidia.com/compute/cuda/repos/ubuntu1604/x86_64/7fa2af80.pub
sudo add-apt-repository "deb https://developer.download.nvidia.com/compute/cuda/repos/ubuntu1604/x86_64/ /"
sudo apt-get update
wget https://developer.download.nvidia.com/compute/machine-learning/repos/ubuntu1604/x86_64/nvidia-machine-learning-repo-ubuntu1604_1.0.0-1_amd64.deb
sudo apt install ./nvidia-machine-learning-repo-ubuntu1604_1.0.0-1_amd64.deb
sudo apt-get update
wget https://developer.download.nvidia.com/compute/machine-learning/repos/ubuntu1604/x86_64/libnvinfer7_7.1.3-1+cuda11.0_amd64.deb
sudo apt install ./libnvinfer7_7.1.3-1+cuda11.0_amd64.deb
sudo apt-get update

# Install development and runtime libraries (~4GB)
sudo apt-get install --no-install-recommends \
    cuda-11-0 \
    libcudnn8=8.0.4.30-1+cuda11.0  \
    libcudnn8-dev=8.0.4.30-1+cuda11.0


# Reboot. Check that GPUs are visible using the command: nvidia-smi

# Install TensorRT. Requires that libcudnn7 is installed above.
sudo apt-get install -y --no-install-recommends \
    libnvinfer7=7.1.3-1+cuda11.0 \
    libnvinfer-dev=7.1.3-1+cuda11.0 \
    libnvinfer-plugin7=7.1.3-1+cuda11.0 \
    libnvinfer-plugin-dev=7.1.3-1+cuda11.0

Configuración en Windows

Consulta los requisitos de hardware y los requisitos de software mencionados más arriba. Lee la guía de instalación de CUDA® para Windows.

Asegúrate de que los paquetes de software de NVIDIA instalados coincidan con las versiones que se muestran más arriba. En particular, TensorFlow no se cargará sin el archivo cuDNN64_8.dll. Para usar otra versión, consulta la guía de compilación con el código fuente en Windows.

Agrega los directorios de instalación CUDA®, CUPTI y cuDNN a la variable de entorno %PATH%. Por ejemplo, si instalas el kit de herramientas CUDA® en C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v11.0 y cuDNN en C:\tools\cuda, actualiza tu %PATH% para que coincida:

SET PATH=C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v11.0\bin;%PATH%
SET PATH=C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v11.0\extras\CUPTI\lib64;%PATH%
SET PATH=C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v11.0\include;%PATH%
SET PATH=C:\tools\cuda\bin;%PATH%