Questo tutorial spiega come usare il tool di generazione di immagini StyleGAN per mezzo dell’ambiente di sviluppo Google Colab, dotato di accelerazione GPU e TPU gratuita. Colab, abbreviazione di Colaboratory, è messo a disposizione da parte di Google per piccole ricerche che richiedano hardware non sempre alla portata di tutti (Tesla K80/ Tesla T4). Colab può essere usato per 12 ore, al termine delle quali il runtime si disconnette e i dati vengono cancellati.
Allenare StyleGAN richiede giorni di training, e cioè necessita più sessioni di Colab. Bisognerà quindi, prima dello scadere delle 12 ore del runtime, salvare i dati per poter poi far ripartire il calcolo dal punto in cui lo avevano precedentemente fermato.
Detto ciò, il primo passo da fare è quello di collegarci a Google Colab:
https://colab.research.google.com
Selezioniamo File>New Python 3 Notebook dal menù in cima alla pagina:
Abbiamo quindi creato il notebook su cui far partire StyleGAN. Si dovrebbe presentare come in immagine:
StyleGAN richiede accelerazione GPU, cioè una scheda grafica su cui distribuire numerose operazioni di calcolo. Selezioniamo quindi dal menù in cima alla pagina Edit>Notebook settings.
Si dovrebbe aprire una finestra come questa, dalla quale selezioniamo la voce GPU come acceleratore hardware. Clicchiamo su ‘Save’ per salvare l’opzione selezionata.
A questo punto tutto è pronto per installare StyleGAN. Scriviamo quindi il codice seguente nella cella di codice (rettangolo grigio) e clicchiamo il tasto play. Questo passo provvede a clonare il programma da Github.
!git clone https://github.com/NVlabs/stylegan.git
Passiamo ora al dataset, cioè all’insieme di immagini che forniremo all’algoritmo per imparare a generarne altre simili.
Carichiamo quindi tutte le nostre immagini su Google Drive, all’interno di una cartella che chiamiamo “dataset”. Per accedere alle immagini su Google Colab, dobbiamo montare collegare Google Drive, cioè importare tutti i contenuti all’interno del nostro Colab. Clicchiamo quindi su “Mount drive” nella sidebar laterale di Colab, come in immagine:
Apparirà quindi una cella di codice nello spazio di lavoro. Eseguiamola. Una volta terminato avremo importato i contenuti di Google Drive, i quali saranno visibili
nella sezione “Files” della sidebar a sinistra, all’interno della cartella “drive”.
StyleGAN richiede che tutte le immagini siano dello stesso formato e abbiano le stesse dimensioni. Per assicurarci di questo installiamo Augmentor, il quale provvederà a ritagliare tutte le foto nella dimensione 512×512, o, in alternativa, a quella che specificherete.
Prima di farlo, creiamo però una cartella dove contenere le immagini ritagliate. Copiate ed eseguite questo codice in una nuova cella. Esso provvederà a creare una cartella apposita all’interno di Stylegan.
import os
if not os.path.exists('/content/stylegan/immagini'):
os.makedirs('/content/stylegan/immagini')
Passiamo quindi all’installazione di Augmentor. Apriamo una nuova cella e incolliamo questo codice. Al posto di “path/to/images” inserite l’indirizzo della cartella “dataset” che abbiamo importato da Google Drive (dovrebbe apparire più o meno così: “/content/drive/My Drive/dataset“). L’altro indirizzo da specificare è quello di destinazione delle immagini, ovvero come precedemente specificato “/content/stylegan/immagini”.
!pip install Augmentor
p = Augmentor.Pipeline("/path/to/images", "/content/stylegan/immagini", save_format="JPEG")
p.resize(probability=1.0, width=512, height=512)
p.random_color(probability=1.0,min_factor=0.5,max_factor=0.9)
p.process()
p.sample(3000)
In particolare il comando p.sample(numero) stabilisce a quante immagini allargare il dataset. Nel machine learning infatti servono molti elementi con cui effettuare il training e quindi, nel caso le immagini a disposizione non fossero di numero sufficiente (almeno 1000), Augmentor ne creerà 3000, modificando piccoli parametri come rotazione e colore ad ogni singola foto in modo da ottenere da ogni singola immagine più copie con cui far allenare StyleGAN.
Una volta fatto ciò indichiamo al programma la cartella dove si trovano le immagini, affinché lui le configuri come dataset. Il codice seguente fa partire il file dataset_tool.py, e crea un nuovo dataset chiamato smalls a partire dalla cartella immagini.
!python /content/stylegan/dataset_tool.py create_from_images datasets/smalls/ /content/stylegan/immagini/
A questo punto, va fatta un’ultima modifica al file train.py per impostare alcuni parametri, ovvero specificare quale dataset di immagini utilizzeremo (un data “custom”, cioè personalizzato, che abbiamo chiamato smalls nel passaggio precedente), e quante GPU utilizzeremo nel calcolo (in questo caso 1, cioè quella che Colab ci mette a disposizione).
Colab non permette di modificare i file tramite un editor, quindi sarà necessario scaricare train.py dalla cartella stylegan, quindi cancellarlo da Colab una volta scaricato (tasto destro del mouse sul file e clic su “Delete”), e poi effettuare l’upload della versione modificata.
Per settare un parametro va rimosso il cancelletto # di fronte alla riga corrispondente. Le righe,già modificate secondo i nostri parametri, sono quelle mostrate nell’immagine seguente.
Per spostare i file all’interno di Colab, può ritornare utile usare questo codice in python:
import shutil
shutil.move('/content/cartelladovesitrovailfile', '/content/cartelladovespostareilfile')
Una volta modificato il file train.py, è possibile fare partire il training vero e proprio:
!python /content/stylegan/train.py
StyleGAN provvederà a salvare una copia del lavoro svolto ogni 10 tick, ovvero all’incirca ogni 1000 iterazioni. La copia viene salvata con il nome di network_snapshot_numero dell’iterazione corrispondente.pkl
tick 1 kimg 140.3 lod 7.00 minibatch 256 time 16m 19s sec/tick 418.8 sec/kimg 2.99 maintenance 560.3 gpumem 3.8
network-snapshot-000140 time 4m 57s fid50k 411.7487
tick 2 kimg 280.6 lod 7.00 minibatch 256 time 30m 44s sec/tick 348.9 sec/kimg 2.49 maintenance 515.8 gpumem 3.8
tick 3 kimg 420.9 lod 7.00 minibatch 256 time 36m 34s sec/tick 347.7 sec/kimg 2.48 maintenance 2.8 gpumem 3.8
tick 4 kimg 561.2 lod 7.00 minibatch 256 time 42m 26s sec/tick 348.7 sec/kimg 2.49 maintenance 3.3 gpumem 3.8
tick 5 kimg 681.5 lod 6.87 minibatch 128 time 52m 00s sec/tick 570.3 sec/kimg 4.74 maintenance 3.0 gpumem 4.2
tick 6 kimg 801.8 lod 6.66 minibatch 128 time 1h 03m 16s sec/tick 673.2 sec/kimg 5.60 maintenance 3.4 gpumem 4.2
tick 7 kimg 922.1 lod 6.46 minibatch 128 time 1h 14m 37s sec/tick 677.9 sec/kimg 5.63 maintenance 2.8 gpumem 4.2
tick 8 kimg 1042.4 lod 6.26 minibatch 128 time 1h 25m 58s sec/tick 677.5 sec/kimg 5.63 maintenance 3.1 gpumem 4.2
tick 9 kimg 1162.8 lod 6.06 minibatch 128 time 1h 37m 22s sec/tick 681.0 sec/kimg 5.66 maintenance 3.1 gpumem 4.2
tick 10 kimg 1283.1 lod 6.00 minibatch 128 time 1h 48m 37s sec/tick 672.1 sec/kimg 5.59 maintenance 2.8 gpumem 4.2
network-snapshot-001283 time 5m 03s fid50k 331.6563
È necessario ricordarsi di salvare l’ultimo snapshot prima dello scadere delle 12 ore concesse da Google Colab. La creazione e il download del file è molto lenta (circa 300 MB di peso) e possono essere necessarie alcune decine di minuti.
Per riprendere il training al punto in cui si era rimasti, creiamo una nuova sessione su Colab, cliccando su Runtime > Reset all runtimes dal menù in cima alla pagina. Ricordiamoci di selezionare nuovamente un’accelerazione GPU per il nostro lavoro. Carichiamo il file pkl precedentemente scaricato e, e rieseguiamo nuovamente ogni punto di questo tutorial.
Prima di avviare il training nuovamente sarà però necessario modificare il file training_loop.py all’interno della cartella options di StyleGAN, specificando dove si trova il file pkl, e da quale iterazione riprendere, come mostrato nell’immagine successiva.
