TensorFlow com GPU no Windows usando o TensorFlow-DirectML-Plugin

Com o fim do suporte a GPU no Windows do TensorFlow, resolvi dar uma olhada nas alternativas. Não tive sucesso com o WSL2, pois não consegui fazer o CUDA funcionar corretamente nele. Mas o TensorFlow-DirectML-Plugin funcionou! Infelizmente a versão dele que está no PyPi suporta apenas o TensorFlow 2.10, ou seja, a mesma versão que ainda tem suporte nativo à GPUs no Windows.

De qualquer forma devo permanecer com o TensorFlow-DirectML-Plugin em que pelo menos há esperança de atualizações para as novas versões do TensorFlow no futuro, o que não acontecerá com a versão nativa.

Também achei mais simples a alternativa de instalar apenas o plugin e ele já instalar a versão suportada do TensorFlow. Combinar as versões corretas de TensorFlow, CUDA Toolkit e cuDNN era bem chato e o Anaconda não ajudava em nada, instalando versões incompatíveis.

A “receita” que estou utilizando com o Anaconda está a seguir. A instalação do TensorFlow-DirectML-Plugin é  com o pip pois ele não está disponível nos repositórios do Anaconda:

conda create --name dml python=3.10
conda activate dml
pip install tensorflow-directml-plugin
conda install pandas scikit-learn pillow spyder

A instalação do plugin também já faz automaticamente a instalação do TensorFlow na versão 2.10. Não adianta atualiza-lo manualmente depois pois não vai funcionar. Não adianta instalar outra versão mais nova, pois com a instalação do plugin ele faz o downgrade. Estou usando o Python 3.10 porque na versão 3.11 o plugin não instala (não tem suporte da versão que está no PyPi)

Cabe destacar também que estou instalando os demais pacotes e IDE que estou usando no meu projeto atual (Pandas, Pillow, Scikit-Learn, Spyder) e instalando-os com o conda.

Para testar se a GPU está funcionando eu uso o seguinte comando:

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