Installare TensorFlow GPU su Ubuntu in WSL 2

Scrivo questa breve guida come appunto per il futuro, dopo che per una giornata sono impazzito dietro all'installazione di TensorFlow che sfruttasse la GPU su Ubuntu in WSL2.

La pubblico qui in caso possa essere utile a qualcun'altro nella mia stessa situazione.

Procedendo con l'installazione classica infatti, nonostante avessi seguito comunque tutto il processo che trovavo sulla guida ufficiale tensorflow (probabilmente outdated), dopo la creazione di un semplice modello ottenevo questo errore da parte di tensorflow:

2019-11-14 22:27:40.533381: W tensorflow/core/common_runtime/gpu/gpu_device.cc:1641] 
Cannot dlopen some GPU libraries. 
Please make sure the missing libraries mentioned above are installed properly if you would like to use GPU. 
Follow the guide at https://www.tensorflow.org/install/gpu for how to download and setup the required libraries for your platform.
Skipping registering GPU devices...

Alla fine, cercando un po' online e tentando e ritentando, sono riuscito a risolvere e ora tensorflow riesce ad accedere alla mia 4090 su Ubuntu in WSL 2. Ecco qui i passaggi dal primo all'ultimo.

Nota importante: lato mio Windows è già correttamente configurato per lavorare con CUDA e con la GPU. Lo scopo di questa guida è far funzionare tutto su WSL oltre che su Windows.

Prima cosa installiamo Ubuntu su WSL

Apriamo Power Shell e digitiamo:

wsl --install -d Ubuntu

Volendo possiamo anche sceglere la distrubuzione tra quelle disponibili digitando

wsl --list -o

Dopo aver selezionato username e password, troveremo sul nostro sistema operativo Ubuntu pronto all'uso.

A questo punto, possiamo confermare la versione della distribuzione installata digitando:

lsb_release -a

Nel mio caso ottengo questo:

No LSB modules are available.
Distributor ID: Ubuntu
Description: Ubuntu 22.04.3 LTS
Release: 22.04
Codename: jammy

Ma ovviamente ogni caso è diverso, in base alla distribuzione installata.

Fatto questo, facciamo subito un aggiornamento di tutto il sistema con:

sudo apt update
sudo apt upgrade

Installiamo tutta l'infrastruttura necessaria

Ora procediamo installando tensorflow, ma prima di questo, è importante strutturare tutto il necessario, in modo tale che possiamo installare tensorflow su un ambiente virtuale ad hoc che crearemo con virtualenv.

Andiamo quindi ad installare innanzitutto pip, digitando

sudo apt install python3-pip

E a questo punto possiamo installare VirtualEnv, digitando:

pip install virtualenv

A questo punto potreste otterrere questo WARNING:

WARNING: The script virtualenv is installed in '/home/simone/.local/bin' which is not on PATH.
Consider adding this directory to PATH or, if you prefer to suppress this warning, use --no-warn-script-location.

Qui potete scegliere se aggiungere la directory in PATH, tramite:

export PATH="$PATH:/home/simone/.local/bin"

E in questo caso potrete poi eseguire virtualenv semplicemente con:

virtualevn NOMEAMBIENTE

Oppure se non aggiungere la directory in PATH e poi eseguire virtualenv con:

python3 -m virtualevn NOMEAMBIENTE

Fatto questo, andomo creare un ambiente virtuale che chiamiamo venv (ma possiamo chiamarlo come vogliamo) e poi lo apriamo. Nota che qui sarebbe il caso di fare prima di tutto una cartella ad hoc all'interno della quale tenere tutti i file del progetto e in cui creare l'ambiente. In questo caso, dato che distruggerò a WSL dopo la guida, terrò tutto in home.

Una volta creato l'ambiente (leggi i comandi sopra), otterremo qualcosa di questo tipo:

created virtual environment CPython3.10.12.final.0-64 in 229ms
creator CPython3Posix(dest=/home/simone/venv, clear=False, no_vcs_ignore=False, global=False)
seeder FromAppData(download=False, pip=bundle, setuptools=bundle, wheel=bundle, via=copy, app_data_dir=/home/simone/.local/share/virtualenv)
added seed packages: pip==24.0, setuptools==69.1.0, wheel==0.42.0
activators BashActivator,CShellActivator,FishActivator,NushellActivator,PowerShellActivator,PythonActivator

E lo possiamo aprire così (ovviamente l'istruzione cambia se avete scelto un diverso nome per l'ambiente virtuale):

source venv/bin/activate

La conferma dell'attivazione ce l'avremmo dalla comparsa di una string che contiene il nome dell'ambiente tra parantesi prima del nostro nome utente sulla shell.

Passando quindi (nel mio caso) da questo:

simone@i9-64gb-RTX4090:~$

A questo

(venv) simone@i9-64gb-RTX4090:~$

A questo punto possiamo procedere con l'installazione di tensorflow e andiamo ad eseguire:

pip install tensorflow

Procediamo con l'installazione di CUDA e cuDNN

Cosa sono CUDA e cuDNN?

Ora dobbiamo installare CUDA e cuDNN, ma per farlo dobbiamo prima di tutto capire quali sono le versioni di cui abbiamo bisogno!

Prima però facciamo chiarezza su che cosa sono e a che cosa servono!

Le CUDA (Compute Unified Device Architecture) sono un'architettura di calcolo parallelo sviluppata da NVIDIA. Consentono di sfruttare la potenza di calcolo delle GPU per eseguire operazioni computazionali ad alte prestazioni. Le GPU sono particolarmente adatte per compiti che richiedono un alto grado di parallelismo, come il deep learning, la simulazione scientifica e il rendering grafico.

CuDNN (CUDA Deep Neural Network library) invece, è una libreria sviluppata da NVIDIA per accelerare le operazioni di deep learning sulle GPU NVIDIA. CuDNN fornisce implementazioni ottimizzate di algoritmi comuni utilizzati nelle reti neurali profonde, come la convoluzione, la normalizzazione dei batch e la propagazione all'indietro. Questo permette di eseguire addestramenti e inferenze neurali più velocemente sfruttando la potenza di calcolo delle GPU.

Quindi CUDA è l'architettura di base che consente di utilizzare le GPU per il calcolo parallelo, mentre CuDNN è una libreria ottimizzata per eseguire operazioni di deep learning su queste GPU. Utilizzando entrambi, è possibile ottenere prestazioni ottimali nell'addestramento e nell'esecuzione di reti neurali profonde.

Procediamo con l'installazione di CUDA e cuDNN

Prima di tutto dobbiamo capire che versione di tensorflow abbiamo installato sul nostro sistema.

Per farlo, dopo aver abilitato il nostro ambiente virtuale (step precedente), eseguiamo:

python3 -c "import tensorflow as tf; print(tf.__version__)"

Questo, oltre ad una probabile lista di warning vari, alla fine dell'output ci darà la versione di tensorflow che abbiamo installato.

Ad esempio io ottengo: 2.16.1

Dunque so che la versione di tensorflow installata è la 2.16.1

E questo è molto importante perché proprio in base a quale è la versione di tensorflow installata io dovrò scegliere le versioni di CUDA e di cuDNN da installare.

Per fare questo, utilizziamo questo link: https://www.tensorflow.org/install/source#gpu

Qui troveremo una tabella che per ogni versione di tensorflow ci da una serie di informazioni, a noi interessano 3 info

Version
Python version
cuDNN
CUDA

Nel mio caso specifico (che probabilimente sarà diverso di volta in volta in base al periodo in cui si legge questa guida), la situzione è questa:

Version: tensorflow-2.16.1
Python version: 3.9-3.12
cuDNN: 8.9
CUDA: 12.3

Dunque so che per la versione di tensorflow tensorflow-2.16.1 (che è quella che io ho installato), ho bisogno di avere:

Una versione di python che va dalla 3.9 alla 3.12
cuDNN versione 8.9
CUDA 12.3

La versione di python, generalmente dovrebbe essere OK, ma possiamo comunque verificare digitando:

python --version

Nel mio caso ottengo: Python 3.10.12

Il che è in linea con le specifiche.

Per quanto riguarda invece CUDA e cuDNN dobiamo installare tutto manualmente: incominciamo dall'installazione di CUDA.

Installazione di CUDA

Una volta identificata la versione di CUDA di cui abbiamo bisogno, andiamo a questa pagina e clicchiamo sul link relativo alla versione più aggiornata della versione di cui abbiamo bisogno: https://developer.nvidia.com/cuda-toolkit-archive

Ad esempio, ne mio caso trovo diverse versioni, ma l'ultima relativa alla 12.3 (di cui ho bisogno) è la CUDA Toolkit 12.3.2 di gennaio 2024.

Oltre a questa ci sono release precedenti della 12.3 (ad esempio la CUDA Toolkit 12.3.1 di novembre 2023 e altre ancora),versioni precedenti alla 12.3 (ad esempio 12.2.x, 12.1.x, e così via) e versioni successive alla 12.3 (nel mio caso ad esempio c'è la 12.4.0 di marzo 2024).

Poiché la versione di cui io ho bisogno è la 12.3 e la più aggiornata è la 12.3.2, allora cliccherò su questo link.

Una volta cliccato sul link si aprirà un wizard che mi guiderà alla scelta della versione corretta.

Nel mio caso sceglierò le seguenti impostazioni:

Linux
x86_64
WSL-Ubuntu
2.0
deb (local)

E una volta fatte tutte queste scelte avremmo a disposizione i comandi da copia-incollare sulla nostra shell.

Nel mio caso ottengo questi (ma attenzione ogni caso è sé stante, utilizzare il comando ottenuto dalla pagina dopo il setup sul wizard):

wget https://developer.download.nvidia.com/compute/cuda/repos/wsl-ubuntu/x86_64/cuda-wsl-ubuntu.pin
sudo mv cuda-wsl-ubuntu.pin /etc/apt/preferences.d/cuda-repository-pin-600
wget https://developer.download.nvidia.com/compute/cuda/12.3.1/local_installers/cuda-repo-wsl-ubuntu-12-3-local_12.3.1-1_amd64.deb
sudo dpkg -i cuda-repo-wsl-ubuntu-12-3-local_12.3.1-1_amd64.deb
sudo cp /var/cuda-repo-wsl-ubuntu-12-3-local/cuda-*-keyring.gpg /usr/share/keyrings/
sudo apt-get update
sudo apt-get -y install cuda-toolkit-12-3

Nota che questo processo di installazione potrebbe richedere un po' di tempo, in base alla propria connessione.

Nel mio caso con la connessione che ho disponibile ora qui nelle Azzorre, ci ha messo una quarantina di minuti, ma in altri test è riuscito a farcela anche in una quindicina. Ma con una connessione migliore di quella che ho io in questi giorni non ci metterete troppo 🙂

Ok fatto questo, esegendo

nvcc -V

Dovremmo ottenere un errore di questo tipo:

Command 'nvcc' not found, but can be installed with: sudo apt install nvidia-cuda-toolkit

Per risolvere il problema e rendere effettiva l'installazione di CUDA aggiungiamo un paio di righe al file .bashrc

Per fare questo apriamo il file con nano digitando:

nano ~/.bashrc

Ci spostiamo alla fine del file (con le frecce della tastiera e non con il mouse) e, una volta raggiunta la fine aggiungiamo queste righe:

export PATH=/usr/local/cuda/bin${PATH:+:${PATH}}
export LD_LIBRARY_PATH=/usr/local/cuda-VERSION/lib64\
                         ${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}}

IMPORTANTE: al posto di VERSION in cuda-VERSION, dobbiamo inserire la versione di CUDA che abbiamo installato. Ad esempio, nel mio caso è la 12.3

Dunque, nel mio caso avremmo:

export PATH=/usr/local/cuda/bin${PATH:+:${PATH}}
export LD_LIBRARY_PATH=/usr/local/cuda-12.3/lib64\
                         ${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}}

E per salvare premiamo CTRL+X, poi Y e poi INVIO

A questo punto aggiornamo il file digitando

. ~/.bashrc

(nota che questo disattiverà l'ambiente virtuale che abbiamo attivato, ma per ora non è un problema)

A questo punto, digitando:

nvcc -V

se tutto è andato bene, otterremmo qualcosa di questo tipo:

nvcc: NVIDIA (R) Cuda compiler driver
Copyright (c) 2005-2023 NVIDIA Corporation
Built on Fri_Nov__3_17:16:49_PDT_2023
Cuda compilation tools, release 12.3, V12.3.103
Build cuda_12.3.r12.3/compiler.33492891_0

E leggendo bene, questo ci conferma che CUDA V12.3.103 è installato correttamene.

A questo punto, passiamo all'installazione dei cuDNN

Installazione di cuDNN

OK, una volta installato CUDA, procediamo all'installazione di cuDNN. Per fare questo dobbiamo fare una procedura simile alla precedente.

Prima di tutto andiamo a questo link e clicchiamo sul link relativo alla versione più aggiornata della versione di cui abbiamo bisogno compatibile con la nostra versione di CUDA: https://developer.nvidia.com/rdp/cudnn-archive

Ad esempio, nel mio caso trovo diverse versioni, ma l'ultima relativa alla 8.9 (di cui ho bisogno, vedi sopra) e compatibile con CUDA 12.x è la 8.9.7 del 5 dicembre 2023.

Cliccando su quel link, nel mio caso non parte alcun wizard, ma mi scarica direttamente il file su windows. Allora una volta scaricato il file, lo copio-incollo direttamente sull'installazione di Ubuntu attraverso il File Explorer sotto Ubuntu-22.04 > Home > Username (ovviamente ogni caso potrebbe essere differente)

Fatto questo procediamo all'installazione di cuDNN con:

sudo apt install ./FILANAME.DEB

Nel mio caso, il nome del file è:

cudnn-local-repo-ubuntu2204-8.9.7.29_1.0-1_amd64

E dunque procedo all'installazione digitando:

sudo apt install ./cudnn-local-repo-ubuntu2204-8.9.7.29_1.0-1_amd64.deb

A questo punto, nel mio caso, ottengo questo messaggio (ma potrebbe essere diverso di caso in caso):

The public cudnn-local-repo-ubuntu2204-8.9.7.29 GPG key does not appear to be installed.
To install the key, run this command:
sudo cp /var/cudnn-local-repo-ubuntu2204-8.9.7.29/cudnn-local-08A7D361-keyring.gpg /usr/share/keyrings/

N: Download is performed unsandboxed as root as file '/home/simone/cudnn-local-repo-ubuntu2204-8.9.7.29_1.0-1_amd64.deb' couldn't be accessed by user '_apt'. - pkgAcquire::Run (13: Permission denied)

Facciamo come dice e digitiamo:

sudo cp /var/cudnn-local-repo-ubuntu2204-8.9.7.29/cudnn-local-08A7D361-keyring.gpg /usr/share/keyrings/

Poi procediamo digitando nuovamente:

sudo apt install ./cudnn-local-repo-ubuntu2204-8.9.7.29_1.0-1_amd64.deb

E otteniamo:

Reading package lists… Done
Building dependency tree… Done
Reading state information… Done
Note, selecting 'cudnn-local-repo-ubuntu2204-8.9.7.29' instead of './cudnn-local-repo-ubuntu2204-8.9.7.29_1.0-1_amd64.deb'
cudnn-local-repo-ubuntu2204-8.9.7.29 is already the newest version (1.0-1).
0 upgraded, 0 newly installed, 0 to remove and 2 not upgraded.

E questo punto, il tocco finale

Installiamo manualmente libcudnn

libcudnn è una libreria che fornisce un insieme di funzioni ottimizzate per l'accelerazione di operazioni comuni nelle reti neurali profonde utilizzando le GPU NVIDIA. Semplifica e velocizza il processo di addestramento e inferenza delle reti neurali profonde utilizzando le GPU NVIDIA.

Per farlo, nel mio caso (che ho le cuDNN versione 8.x) digito:

sudo apt update
sudo apt install libcudnn8
sudo apt install libcudnn8-dev
sudo apt install libcudnn8-samples

Confermiamo che tutto funziona correttamente!

Siamo pronti ora a confermare che tutto è stato impostato correttamente. Per farlo riattiviamo il virtual envirorment (leggi sopra), e una volta che abbiamo il virtual enviromente attivo (è importante), digitiamo:

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

Se tutto è andato bene, otterremo un output come questo:

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

A conferma del fatto che la GPU è stata caricata correttamente ed è disponibile per essere utilizzata da TensorFlow

Risorse utili & Credits

Sono impazzito una mezza giornata a capire perchè non funzionasse e come risolvere il problema. Cercando in rete ho trovato molte risorse, alcune utili, altre meno.

Vi lacio qui quelle che mi hanno aiutato di più!

Immagine di copertina, generata da e disponibile su Civitai.com all'indirizzo: https://civitai.com/images/4839813