Instalando o Tensorflow com suporte a GPU no Windows com WSL2

O TensorFlow é uma biblioteca de aprendizado de máquina de código aberto desenvolvida pelo Google. É uma ferramenta poderosa para a construção de modelos de aprendizado de máquina, mas alguns desenvolvimentos se beneficiam quando utilizamos uma GPU ao invés da CPU, um bom exemplo são os modelos de visão computacional.

Entretanto, para aqueles que utilizam o sistema Windows, é preciso seguir alguns passos para conseguir usufruir de uma placa de vídeo. Isso porque a última versão do TensorFlow com suporte nativo para GPU no Windows foi a 2.10, que deixou de estar disponível para instalação via pip. Para utilizar a GPU, é necessário utilizar a biblioteca no WSL2 ou instalar o TensorFlow-DirectML-Plugin, que também é válido para quem possui placas AMD.

Devo ressaltar que a documentação original encontra-se no seguintes links:

Instale o TensorFlow com pip

Suporte a GPUs | TensorFlow

Antes de qualquer coisa, é importante verificar os requisitos do sistema, conforme a documentação:

• Placa GPU NVIDIA® com arquiteturas CUDA® 3.5, 5.0, 6.0, 7.0, 7.5, 8.0 e superior (https://developer.nvidia.com/cuda-gpus)
• Para GPUs com arquiteturas CUDA® não suportadas, ou para evitar a compilação JIT do PTX, ou para usar versões diferentes das bibliotecas NVIDIA®, consulte o guia Linux build from source .
• Os pacotes não contêm código PTX, exceto para a arquitetura CUDA® suportada mais recente; portanto, o TensorFlow falha ao carregar em GPUs mais antigas quando CUDA_FORCE_PTX_JIT=1 está definido. (Consulte Compatibilidade de aplicativos para obter detalhes.)

Sistemas Operacionais aceitos:

• Ubuntu 16.04 ou superior (64 bits)
• macOS 10.12.6 (Sierra) ou superior (64 bits) (sem suporte de GPU)
• Windows Native — Windows 7 ou superior (64 bits) (sem suporte de GPU após TF 2.10)
• Windows WSL2 — Windows 10 19044 ou superior (64 bits)

Requisitos de software:

• Python 3.9–3.11
• pip versão 19.0 ou superior para Linux (requer suporte para manylinux2014 ) e Windows. pip versão 20.3 ou superior para macOS.
• Windows Native requer Microsoft Visual C++ Redistributable para Visual Studio 2015, 2017 e 2019

Os seguintes softwares NVIDIA® são necessários apenas para suporte de GPU.

• Drivers de GPU NVIDIA® versão 450.80.02 ou superior.
• Kit de ferramentas CUDA® 11.8 .
• cuDNNSDK 8.6.0 .

Instalação:

O primeiro passo é a instalação do WSL2, para isso abra o PowerShell e insira o seguinte comando:

wsl --install

No final da instalação, será necessário reiniciar o computador para habilitar os recursos necessários para executar o WSL.

Você pode usar o seguinte comando para listar as distribuições Linux disponíveis para download na loja online:

wsl --list --online

Uma vez escolhida a distribuição, basta executar o comando de instalação, conforme o exemplo:

wsl.exe --install <Distro>

Após a instalação, você irá configurar suas informações de usuário do Linux.

Agora iremos configurar a GPU, para isso é preciso que você instale o driver de sua placa caso ainda não tenha instalado.

O comando a seguir irá te ajudar a verificar se o driver esta instalado:

nvidia-smi

Por fim, instalaremos o TensorFlow. Para isso, é recomendado instalar e/ou atualizar o pip:

sudo apt update

sudo apt install python3-pip

pip install --upgrade pip

Caso deseje, você pode instalar o TensorFlow em um ambiente virtual do Python ou Anaconda. Aqui, estarei exemplificando com o uso do virtualenv.

sudo apt install python3-virtualenv

virtualenv venv 

source venv/bin/activate

Agora com o ambiente iniciado, vamos passar para a instalação do tensorflow:

 pip install tensorflow[and-cuda]

Pronto! Estamos aptos a testar a instalação agora. Para isso, execute:

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

A saida será algo semelhante a essa imagem:

Também podemos executar o seguinte comando para verificar a configuração da GPU:

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

A saída deve ser uma lista com os dispositivos GPU.

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

Bônus: Utilizando o JupyterLab

Uma das maneiras mais simples de utilizar o seu TensorFlow configurado para utilizar a sua placa de vídeo é através da instalação do JupyterLab a partir do sistema Linux instalado no WSL com o conforto do Windows e o seu navegador favorito.

Para isso, basta instalar o JupyterLab com os seguintes comandos:

pip install jupyterlab

jupyter lab

Tudo pronto! Agora, basta acessar o endereço que será disponibilizado.

Por fim para desabilitar o WSL é só executar o seguinte comando:

exit
wsl --shutdown

E é isso pessoal, espero ter ajudado !!

Até a próxima!