stylegan2 pytorch

Einfacher Stylegan2 für Pytorch

Einfache Pytorch-Implementierung von Stylegan2 basierend auf https://arxiv.org/abs/1912.04958, die vollständig aus der Befehlszeile trainiert werden kann, ohne Codierung erforderlich.

Im Folgenden sind einige Blumen aufgeführt, die nicht existieren.

Auch diese Hände auch nicht

Noch diese Städte

Noch diese Prominenten (ausgebildet von @yoniker)

Installieren

Sie benötigen eine Maschine mit einer GPU und CUDA. Dann installieren Sie das Paket wie dieses Paket

$ pip install stylegan2_pytorch

Wenn Sie einen Windows -Computer verwenden, funktionieren die folgenden Befehle Berichten zufolge.

$ conda install pytorch torchvision -c python
$ pip install stylegan2_pytorch

Verwenden

$ stylegan2_pytorch --data /path/to/images

Das war's. Beispielbilder werden in results/default gespeichert und Modelle werden regelmäßig in models/default gespeichert.

Erweiterte Verwendung

Sie können den Namen Ihres Projekts angeben

$ stylegan2_pytorch --data /path/to/images --name my-project-name

Sie können auch den Ort angeben, an dem Intermediate -Ergebnisse und Modellkontrollpunkte mit gespeichert werden sollten

$ stylegan2_pytorch --data /path/to/images --name my-project-name --results_dir /path/to/results/dir --models_dir /path/to/models/dir

Sie können die Netzwerkkapazität (die standardmäßig auf 16 ) erhöhen, um die Erzeugungsergebnisse auf Kosten eines mehr Speichers zu verbessern.

$ stylegan2_pytorch --data /path/to/images --network-capacity 256

Wenn das Training abgeschaltet wird, wird standardmäßig automatisch von der letzten Kontrolldatei fortgesetzt. Wenn Sie mit neuen Einstellungen neu starten möchten, fügen Sie einfach eine new Flagge hinzu

$ stylegan2_pytorch --new --data /path/to/images --name my-project-name --image-size 512 --batch-size 1 --gradient-accumulate-every 16 --network-capacity 10

Sobald Sie das Training fertiggestellt haben, können Sie Bilder aus Ihrem neuesten Checkpoint aus generieren.

$ stylegan2_pytorch  --generate

Ein Video einer Interpolation durch zwei zufällige Punkte im latenten Raum zu generieren.

$ stylegan2_pytorch --generate-interpolation --interpolation-num-steps 100

Um jeden einzelnen Rahmen der Interpolation zu retten

$ stylegan2_pytorch --generate-interpolation --save-frames

Wenn ein früherer Kontrollpunkt einen besseren Generator enthielt (der häufig als Generatoren gegen Ende des Trainings auftritt), können Sie mit einem anderen Flag von einem vorherigen Kontrollpunkt aus laden

$ stylegan2_pytorch --generate --load-from {checkpoint number}

Eine Technik, die sowohl in Stylegan als auch in Biggan verwendet wird, schneidet die latenten Werte so ab, dass ihre Werte dem Mittelwert nahe kommen. Je klein der Kürzungswert, desto besser erscheinen die Proben auf Kosten der Probenvielfalt. Sie können dies mit dem --trunc-psi steuern, wobei die Werte normalerweise zwischen 0.5 und 1 fallen. Es wird als Standard auf 0.75 eingestellt

$ stylegan2_pytorch --generate --trunc-psi 0.5

Multi-GPU-Training

Wenn Sie eine Maschine mit mehreren GPUs haben, bietet das Repository eine Möglichkeit, alle für das Training zu verwenden. Bei mehreren GPUs wird jede Charge gleichmäßig auf die verfügbaren GPUs aufgeteilt. Beispielsweise wird für 2 GPUs mit einer Chargengröße von 32 jeweils 16 Proben angezeigt.

Sie müssen einfach eine Flagge --multi-gpus hinzufügen, für die sich alles andere erledigt. Wenn Sie sich auf bestimmte GPUs beschränken möchten, können Sie die Umgebungsvariable CUDA_VISIBLE_DEVICES verwenden, um zu steuern, welche Geräte verwendet werden können. (Bsp. CUDA_VISIBLE_DEVICES=0,2,3 Nur Geräte 0, 2, 3 sind verfügbar)

$ stylegan2_pytorch --data ./data --multi-gpus --batch-size 32 --gradient-accumulate-every 1

Niedrige Mengen an Schulungsdaten

In der Vergangenheit brauchte Gans viele Daten, um zu lernen, wie man gut generiert. Das Faces -Modell hat als Beispiel 70.000 hochwertige Bilder von Flickr aufgenommen.

Im Monat Mai 2020 konvergierten Forscher auf der ganzen Welt jedoch unabhängig von einer einfachen Technik, um diese Zahl auf nur 1-2K zu reduzieren. Diese einfache Idee war es, alle erzeugten oder realen Bilder differenzend zu erweitern und während des Trainings in den Diskriminator zu gehen.

Wenn man auf eine geringe Wahrscheinlichkeit ausreicht, werden die Augmentationen nicht in die Generationen eintreten.

In der Einstellung niedriger Daten können Sie die Funktion mit einem einfachen Flag verwenden.

 # find a suitable probability between 0. -> 0.7 at maximum
$ stylegan2_pytorch --data ./data --aug-prob 0.25

Standardmäßig sind die verwendeten Augmentationen translation und cutout . Wenn Sie color hinzufügen möchten, können Sie dies mit dem Argument --aug-types tun.

 # make sure there are no spaces between items!
$ stylegan2_pytorch --data ./data --aug-prob 0.25 --aug-types [translation,cutout,color]

Sie können es an eine beliebige Kombination der drei anpassen, die Sie möchten. Der differenzierbare Augmentationscode wurde von hier aus kopiert und leicht modifiziert.

Wann stoppt ich das Training?

So lange wie möglich, bis das kontroverse Spiel zwischen den beiden neuronalen Netzen auseinander fällt (wir nennen diese Divergenz). Standardmäßig wird die Anzahl der Trainingsschritte für 128x128 -Bilder auf 150000 gesetzt, aber Sie möchten sicherlich höher sind, dass diese Zahl höher ist, wenn der GaN bis zum Ende des Trainings nicht abweist oder wenn Sie bei einer höheren Auflösung trainieren.

$ stylegan2_pytorch --data ./data --image-size 512 --num-train-steps 1000000

Aufmerksamkeit

In diesem Rahmen können Sie auch eine effiziente Form der Selbstbekämpfung zu den festgelegten Schichten des Diskriminators (und der symmetrischen Schicht des Generators) hinzufügen, die die Ergebnisse erheblich verbessern. Je mehr Aufmerksamkeit Sie sich leisten können, desto besser!

 # add self attention after the output of layer 1
$ stylegan2_pytorch - - data . / data - - attn - layers 1 

 # add self attention after the output of layers 1 and 2
# do not put a space after the comma in the list!
$ stylegan2_pytorch - - data . / data - - attn - layers [ 1 , 2 ]

Bonus

Training auf transparenten Bildern

$ stylegan2_pytorch --data ./transparent/images/path --transparent

Gedächtnisüberlegungen

Je mehr GPU -Speicher Sie haben, desto größer und besser wird die Bildgenerierung sein. Nvidia empfahl bis zu 16 GB für das Training 1024x1024 Bilder. Wenn Sie weniger als das haben, können Sie ein paar Einstellungen spielen, damit das Modell passt.

$ stylegan2_pytorch --data /path/to/data 
    --batch-size 3 
    --gradient-accumulate-every 5 
    --network-capacity 16

1. Batchgröße-Sie können die batch-size auf 1 verringern, aber Sie sollten den gradient-accumulate-every entsprechend erhöhen, damit das Mini-Batch das Netzwerk nicht zu klein ist. Dies mag für einen Laien verwirrend sein, also werde ich darüber nachdenken, wie ich die Auswahl von gradient-accumulate-every Zukunft automatisieren würde.

2. Netzwerkkapazität - Sie können die neuronale Netzkapazität verringern, um die Speicheranforderungen zu verringern. Beachten Sie nur, dass dies nachgewiesen wurde, dass dies die Erzeugung beeinträchtigt.

Wenn keines dieser Funktionen funktioniert, können Sie sich mit dem 'Leichtgewicht' Gan zufrieden geben, sodass Sie in angemessener Zeit einen Kompromissqualität in größeren Vorsätzen ausbilden können.

Einsatz auf AWS

Im Folgenden finden Sie einige Schritte, die für die Bereitstellung mit Amazon Web Services hilfreich sein können. Um dies zu verwenden, müssen Sie eine von GPU unterstützte EC2-Instanz bereitstellen. Ein geeigneter Instanztyp stammt aus einer P2- oder P3 -Serie. Ich (iboates) probierte eine P2.xlarge (die billigste Option) und es war ziemlich langsam und in der Tat langsamer als die Verwendung von Google Colab. Leistungsstärkere Instanztypen können besser sein, aber sie sind teurer. Sie können hier mehr darüber lesen.

Setup -Schritte einrichten

1. Archivieren Sie Ihre Trainingsdaten und laden Sie sie in einen S3 -Bucket hoch
2. Bereitstellung Ihrer EC2 -Instanz (ich habe ein Ubuntu AMI verwendet)
3. Melden Sie sich in Ihrer EC2 -Instanz über SSH an
4. Installieren Sie den AWS CLI -Client und konfigurieren Sie es:

sudo snap install aws-cli --classic
aws configure

Sie müssen dann Ihre AWS -Zugriffsschlüssel eingeben, die Sie unter AWS -Verwaltungskonsole> Profil> My Security Anmeldeinformationen> Zugriffsschlüssel aus der Verwaltungskonsole abrufen können

Führen Sie dann diese Befehle aus oder legen Sie sie vielleicht in ein Shell -Skript und führen Sie das aus:

mkdir data
curl -O https://bootstrap.pypa.io/get-pip.py
sudo apt-get install python3-distutils
python3 get-pip.py
pip3 install stylegan2_pytorch
export PATH= $PATH :/home/ubuntu/.local/bin
aws s3 sync s3:// < Your bucket name > ~ /data
cd data
tar -xf ../train.tar.gz

Jetzt sollten Sie in der Lage sein, zu trainieren, indem Sie einfach anrufen stylegan2_pytorch [args] .

Anmerkungen:

• Wenn Sie viele Trainingsdaten haben, müssen Sie möglicherweise zusätzlichen Blockspeicher über EBS bereitstellen.
• Möglicherweise müssen Sie Ihre Daten auf mehrere Archive verteilen.
• Sie sollten dies in einem screen ausführen, damit es nicht beendet wird, sobald Sie sich von der SSH -Sitzung abmelden.

Forschung

FID -Scores

Dank des GetEclectic können Sie jetzt den FID -Score regelmäßig berechnen! Wieder mit einem zusätzlichen Argument, wie unten gezeigt, super einfach gemacht.

Installieren Sie zunächst das pytorch_fid -Paket

$ pip install pytorch-fid

Gefolgt von

$ stylegan2_pytorch --data ./data --calculate-fid-every 5000

FID -Ergebnisse werden an ./results/{name}/fid_scores.txt <Name}/fid_scores.txt angemeldet

Codierung

Wenn Sie Bilder programmgesteuert probieren möchten, können Sie dies mit der folgenden einfachen ModelLoader -Klasse tun.

 import torch
from torchvision . utils import save_image
from stylegan2_pytorch import ModelLoader

loader = ModelLoader (
    base_dir = '/path/to/directory' ,   # path to where you invoked the command line tool
    name = 'default'                   # the project name, defaults to 'default'
)

noise   = torch . randn ( 1 , 512 ). cuda () # noise
styles  = loader . noise_to_styles ( noise , trunc_psi = 0.7 )  # pass through mapping network
images  = loader . styles_to_images ( styles ) # call the generator on intermediate style vectors

save_image ( images , './sample.jpg' ) # save your images, or do whatever you desire

Anmelden zum Experiment -Tracker

Um die Verluste bei einem Open -Source -Experiment -Tracker (AIM) zu protokollieren, müssen Sie lediglich eine zusätzliche Flagge wie dies übergeben.

$ stylegan2_pytorch --data ./data --log

Dann müssen Sie sicherstellen, dass Sie Docker installiert haben. Nach den Anweisungen von AIM führen Sie Folgendes in Ihrem Terminal aus.

$ aim up

Öffnen Sie dann Ihren Browser an der Adresse und Sie sollten sehen

Experimental

Top-K-Training für Generator

Ein neues Papier hat Beweise dafür erbracht, dass der Generator durch einfaches Null der Gradientenbeiträge aus Proben, die vom Diskriminator als gefälscht angesehen werden, erheblich besser lernt und einen neuen Stand der Technik erreicht.

$ stylegan2_pytorch - - data . / data - - top - k - training

Gamma ist ein Zerfallsplan, der das Topk langsam von der vollständigen Chargengröße bis zur Zielfraktion von 50% verringert (auch modifizierbarer Hyperparameter).

$ stylegan2_pytorch - - data . / data - - top - k - training - - generate - top - k - frac 0.5 - - generate - top - k - gamma 0.99

Feature Quantisierung

Eine kürzlich durchgeführte Arbeit berichtete verbesserte Ergebnisse, wenn Zwischendarstellungen des Diskriminators vektor quantisiert sind. Obwohl ich keine dramatischen Veränderungen bemerkt habe, habe ich beschlossen, dies als Funktion hinzuzufügen, damit andere Köpfe da draußen untersuchen können. Um zu verwenden, müssen Sie angeben, welche Ebenen Sie Vektor quantisieren möchten. Die Standard -Wörterbuchgröße beträgt 256 und ist ebenfalls einstellbar.

 # feature quantize layers 1 and 2, with a dictionary size of 512 each
# do not put a space after the comma in the list!
$ stylegan2_pytorch - - data . / data - - fq - layers [ 1 , 2 ] - - fq - dict - size 512

Kontrastive Verlust -Regularisierung

Ich habe kontrastives Lernen auf dem Diskriminator (im Schritt mit dem üblichen GaN -Training) ausprobiert und möglicherweise eine verbesserte Stabilität und Qualität der Endergebnisse beobachtet. Sie können diese experimentelle Funktion mit einem einfachen Flag einschalten, wie unten gezeigt.

$ stylegan2_pytorch - - data . / data - - cl - reg

Relativistischer Diskriminatorverlust

Dies wurde im relativistischen Gan -Papier vorgeschlagen, um das Training zu stabilisieren. Ich habe gemischte Ergebnisse erzielt, werde aber die Funktion für diejenigen enthalten, die damit experimentieren möchten.

$ stylegan2_pytorch - - data . / data - - rel - disc - loss

Nicht konstanter 4x4-Block

Standardmäßig stilft der Stylegan -Architektur einen konstanten, erlernten 4x4 -Block, da er zunehmend am 2 -jährigen gestrichen wird. Dies ist ein experimentelles Merkmal, das es so macht, dass der 4x4 -Block stattdessen aus dem Stilvektor w gelernt wird.

$ stylegan2_pytorch - - data . / data - - no - const

Doppelter Kontrastverlust

Eine kürzlich durchgeführte Arbeit hat vorgeschlagen, dass ein neuer kontrastiver Verlust zwischen den realen und gefälschten Protokollen die Qualität gegenüber anderen Arten von Verlusten verbessern kann. (Der Standard in diesem Repository ist Scharnierverlust, und das Papier zeigt eine leichte Verbesserung)

$ stylegan2_pytorch - - data . / data - - dual - contrast - loss 

Alternativen

Stylegan2 + Unet Diskriminator

Ich habe wirklich gute Ergebnisse mit einem Unet Diskriminator erzielt, aber die architektonische Veränderung war zu groß, um als Option in diesem Repository zu passen. Wenn Sie Perfektion anstreben, können Sie es gerne versuchen.

Wenn Sie möchten, dass ich die königliche Behandlung an eine andere Gan -Architektur (Biggan) gibt, wenden Sie sich an meine E -Mail. Freut mich, deine Tonhöhe zu hören.

Anerkennung

Vielen Dank an Matthew Mann für seinen inspirierenden einfachen Port für TensorFlow 2.0

Referenzen

 @article { Karras2019stylegan2 ,
    title   = { Analyzing and Improving the Image Quality of {StyleGAN} } ,
    author  = { Tero Karras and Samuli Laine and Miika Aittala and Janne Hellsten and Jaakko Lehtinen and Timo Aila } ,
    journal = { CoRR } ,
    volume  = { abs/1912.04958 } ,
    year    = { 2019 } ,
}

 @misc { zhao2020feature ,
    title   = { Feature Quantization Improves GAN Training } ,
    author  = { Yang Zhao and Chunyuan Li and Ping Yu and Jianfeng Gao and Changyou Chen } ,
    year    = { 2020 }
}

 @misc { chen2020simple ,
    title   = { A Simple Framework for Contrastive Learning of Visual Representations } ,
    author  = { Ting Chen and Simon Kornblith and Mohammad Norouzi and Geoffrey Hinton } ,
    year    = { 2020 }
}

 @article {,
    title     = { Oxford 102 Flowers } ,
    author    = { Nilsback, M-E. and Zisserman, A., 2008 } ,
    abstract  = { A 102 category dataset consisting of 102 flower categories, commonly occuring in the United Kingdom. Each class consists of 40 to 258 images. The images have large scale, pose and light variations. }
}

 @article { afifi201911k ,
    title   = { 11K Hands: gender recognition and biometric identification using a large dataset of hand images } ,
    author  = { Afifi, Mahmoud } ,
    journal = { Multimedia Tools and Applications }
}

 @misc { zhang2018selfattention ,
    title   = { Self-Attention Generative Adversarial Networks } ,
    author  = { Han Zhang and Ian Goodfellow and Dimitris Metaxas and Augustus Odena } ,
    year    = { 2018 } ,
    eprint  = { 1805.08318 } ,
    archivePrefix = { arXiv }
}

 @article { shen2019efficient ,
    author    = { Zhuoran Shen and
               Mingyuan Zhang and
               Haiyu Zhao and
               Shuai Yi and
               Hongsheng Li } ,
    title     = { Efficient Attention: Attention with Linear Complexities } ,
    journal   = { CoRR } ,  
    year      = { 2018 } ,
    url       = { http://arxiv.org/abs/1812.01243 } ,
}

 @article { zhao2020diffaugment ,
    title   = { Differentiable Augmentation for Data-Efficient GAN Training } ,
    author  = { Zhao, Shengyu and Liu, Zhijian and Lin, Ji and Zhu, Jun-Yan and Han, Song } ,
    journal = { arXiv preprint arXiv:2006.10738 } ,
    year    = { 2020 }
}

 @misc { zhao2020image ,
    title  = { Image Augmentations for GAN Training } ,
    author = { Zhengli Zhao and Zizhao Zhang and Ting Chen and Sameer Singh and Han Zhang } ,
    year   = { 2020 } ,
    eprint = { 2006.02595 } ,
    archivePrefix = { arXiv }
}

 @misc { karras2020training ,
    title   = { Training Generative Adversarial Networks with Limited Data } ,
    author  = { Tero Karras and Miika Aittala and Janne Hellsten and Samuli Laine and Jaakko Lehtinen and Timo Aila } ,
    year    = { 2020 } ,
    eprint  = { 2006.06676 } ,
    archivePrefix = { arXiv } ,
    primaryClass = { cs.CV }
}

 @misc { jolicoeurmartineau2018relativistic ,
    title   = { The relativistic discriminator: a key element missing from standard GAN } ,
    author  = { Alexia Jolicoeur-Martineau } ,
    year    = { 2018 } ,
    eprint  = { 1807.00734 } ,
    archivePrefix = { arXiv } ,
    primaryClass = { cs.LG }
}

 @misc { sinha2020topk ,
    title   = { Top-k Training of GANs: Improving GAN Performance by Throwing Away Bad Samples } ,
    author  = { Samarth Sinha and Zhengli Zhao and Anirudh Goyal and Colin Raffel and Augustus Odena } ,
    year    = { 2020 } ,
    eprint  = { 2002.06224 } ,
    archivePrefix = { arXiv } ,
    primaryClass = { stat.ML }
}

 @misc { yu2021dual ,
    title   = { Dual Contrastive Loss and Attention for GANs } ,
    author  = { Ning Yu and Guilin Liu and Aysegul Dundar and Andrew Tao and Bryan Catanzaro and Larry Davis and Mario Fritz } ,
    year    = { 2021 } ,
    eprint  = { 2103.16748 } ,
    archivePrefix = { arXiv } ,
    primaryClass = { cs.CV }
}