BigGAN PyTorch Download - BigGAN PyTorch Quellcode Download

BigGAN PyTorch

Python

1.0.0

Herunterladen

Biggan-Pytorch

Die offiziell inoffizielle Pytorch -Biggan -Implementierung des Autors.

Dogball? Dogball!

Dieses Repo enthält Code für das 4-8 GPU-Training von Biggans aus großem Umfang GaN-Training für die natürliche Bildsynthese von High Fidelity von Andrew Brock, Jeff Donahue und Karen Simonyan.

Dieser Code stammt von Andy Brock und Alex Andonian.

So verwenden Sie diesen Code

Sie brauchen:

Pytorch, Version 1.0.1
TQDM, Numpy, Scipy und H5Py
Das Bildnetz -Trainingsset

Erstens können Sie optional eine vorverarbeitete HDF5-Version Ihres Zieldatensatzes für schnellere E/A vorbereiten. Im Folgenden (oder nicht) benötigen Sie die Inception -Momente, die zur Berechnung von FID erforderlich sind. Diese können sowohl durch Modifizierung und Laufen erfolgen

sh scripts/utils/prepare_data.sh

Dies wird standardmäßig von Ihrem ImageNet -Trainingssatz angenommen, der in den Root data in diesem Verzeichnis heruntergeladen wird und den zwischengespeicherten HDF5 bei der Auflösung von 128x128 Pixel vorbereitet.

Im Ordner Skripts gibt es mehrere Bash -Skripte, die Biggans mit unterschiedlichen Chargengrößen trainieren. In diesem Code wird davon ausgegangen, dass Sie keinen Zugriff auf einen vollständigen TPU-Pod haben, und entsprechend wird Mega-Batches mithilfe von Gradientenakkumulation (Mittelwertabsolventen über mehrere Minibatches und nur einen Optimiererschritt nach N-Akkumulationen) gefälscht. Standardmäßig trainiert das Drehbuch des Schriftzugs launch_BigGAN_bs256x8.sh ein Biggan-Modell in voller Größe mit einer Stapelgröße von 256 und 8 Gradientenakkumulationen für eine Gesamtgröße von 2048. Auf 8xv100 mit vollem Präzisionstraining (keine Tensor-Kerne) dauert dieses Skript bis zu 150K-Iterationen 15 Tage.

Sie müssen zunächst die maximale Chargengröße finden, die Ihr Setup unterstützen kann. Die hier bereitgestellten vorgeborenen Modelle wurden auf 8xv100 (jeweils 16 GB VRAM) trainiert, was etwas mehr unterstützen kann als der standardmäßige BS256. Sobald Sie dies ermittelt haben, sollten Sie das Skript so ändern, dass die Stapelgröße die Anzahl der Gradientenakkumulationen gleich der gewünschten Gesamtstapelgröße entspricht (Biggan standardmäßig bis 2048).

Beachten Sie auch, dass dieses Skript den --load_in_mem arg verwendet, der die gesamte (~ 64 GB) i128.hdf5 -Datei zum schnelleren Datenladen in RAM lädt. Wenn Sie nicht genug RAM haben, um dies zu unterstützen (wahrscheinlich 96 GB+), entfernen Sie dieses Argument.

Metriken und Probenahme

I believe I can fly!

Während des Trainings gibt dieses Skript Protokolle mit Trainingsmetriken und Testmetriken aus, speichert mehrere Kopien (2 neueste und 5 höchste Punktzahl) der Modellgewichte/Optimiererparameter und produziert jedes Mal, wenn es Gewicht spart. Der Protokollordner enthält Skripte, um diese Protokolle zu verarbeiten und die Ergebnisse mit MATLAB zu zeichnen (sorry nicht sorry).

Nach dem Training kann ein Beispiel für sample.py verwendet werden, um zusätzliche Proben und Interpolationen zu produzieren, mit unterschiedlichen Abschnitten, Stapelgrößen, Anzahl der ständigen STAT -Akkumulationen usw. Das Skript sample_BigGAN_bs256x8.sh für ein Beispiel.

Standardmäßig wird alles in Gewichten/Samples/Protokollen/Datenordnern gespeichert, von denen angenommen wird, dass sie sich im selben Ordner wie dieses Repo befinden. Sie können alle diese mit dem Argument --base_root in einen anderen Basisordner verweisen oder für jeden von diesen mit ihren jeweiligen Argumenten (z. B. --logs_root ) bestimmte Stellen auswählen.

Wir schließen Skripte ein, um Biggan-Deep auszuführen, aber wir haben ein Modell mit ihnen nicht vollständig geschult. Betrachten Sie sie also ungetestet. Zusätzlich schließen wir Skripte ein, um ein Modell auf Cifar auszuführen und SA-Gan (mit EMA) und SN-Gan auf ImageNet auszuführen. Der SA-Gan-Code geht davon aus, dass Sie 4xtitanx (oder gleichwertig in Bezug auf GPU-RAM) haben und mit einer Stapelgröße von 128 und 2 Gradientenakkumulationen ausgeführt werden.

Ein wichtiger Hinweis zu Inception -Metriken

Dieses Repo verwendet das in Pytorch eingebaute Inception-Netzwerk, um IS und FID zu berechnen. Diese Bewertungen unterscheiden sich von den Bewertungen, die Sie mit dem offiziellen TF -Inception -Code erhalten würden, und dienen nur zu Überwachungszwecken! Ausführen von Sample.Py auf Ihrem Modell mit dem Argument --sample_npz und dann inception_tf13 ausführen, um den tatsächlichen Tensorflow IS zu berechnen. Beachten Sie, dass Sie TensorFlow 1.3 oder früher installiert haben müssen, da TF1.4+ das Original Code ist.

Vorbereitete Modelle

Pytorch Inception Score und FID Wir enthalten zwei vorbereitete Modell -Checkpoints (mit G, D, der EMA -Kopie von G, den Optimierern und dem Zustandsdikt):

Der Hauptkontrollpunkt ist für einen Biggan, der bei 128x128 auf ImageNet ausgebildet wurde, unter Verwendung von BS256- und 8 Gradientenakkumulationen, die kurz vor dem Zusammenbruch eingenommen wurden, mit einem TF-Inception-Score von 97,35 +/- 1,79: Link
Ein früherer Kontrollpunkt des ersten Modells (100K G ITERS), bei hoher Leistung, aber kurz vor dem Zusammenbruch, was möglicherweise einfacher zu feinstimmen ist: Link

Vorbereitete Modelle für Places-365 in Kürze.

Dieses Repo enthält auch Skripte für die Portierung des ursprünglichen TFHUB -Biggan -Generators auf Pytorch. Weitere Informationen finden Sie in den Skripten im TFHUB -Ordner.

Feinabstimmung, Verwendung Ihres eigenen Datensatzes oder neue Trainingsfunktionen

That's deep, man

Wenn Sie unterbrochenes Training wieder aufnehmen oder ein vorgebildetes Modell fein abteilen möchten, führen Sie das gleiche Startskript aus, jedoch mit dem hinzugefügten Argument --resume . Experimentennamen werden automatisch aus der Konfiguration generiert, können jedoch mit dem Arg --experiment_name (z.

Um Ihren eigenen Datensatz vorzubereiten, müssen Sie ihn zu DataSets.py hinzufügen und die Convenience -Diktate in utils.py (dset_dict, iMsize_dict, root_dict, nclass_dict, class_per_sheet_dict) für die entsprechenden Metadaten für Ihren Datensatz ändern. Wiederholen Sie den Vorgang in prepe_data.sh (produzieren Sie optional eine vorverarbeitete HDF5 -Kopie und berechnen Sie die Inception -Momente für FID).

Standardmäßig speichert das Trainingsskript die 5 besten Checkpoints, gemessen an der Inception -Punktzahl. Für andere Datensätze als ImageNet kann der Inception -Score ein sehr schlechtes Maß an Qualität sein, sodass Sie wahrscheinlich verwenden möchten --which_best FID .

Um Ihre eigene Trainingsfunktion zu verwenden (z. B. Train a Bigvae): Ändern Sie entweder train_fns.gan_training_function oder fügen Sie einen neuen Zug Fn hinzu und fügen Sie ihn nach der if config['which_train_fn'] == 'GAN': Linie in train.py .

Ordentliches Zeug

Wir enthalten hier die vollständigen Trainings- und Metrikprotokolle zur Referenz. Ich habe festgestellt, dass eines der schwierigsten Dinge beim erneuten Implementieren eines Papiers die Überprüfung sein kann, ob die Protokolle früh im Training ausgerichtet sind, insbesondere wenn das Training mehrere Wochen dauert. Hoffentlich sind diese für zukünftige Arbeiten hilfreich.
Wir fügen eine beschleunigte FID-Berechnung hinzu-die ursprüngliche Scipy-Version kann über 10 Minuten benötigt werden, um die Matrix SQRT zu berechnen. Diese Version verwendet eine beschleunigte Pytorch-Version, um sie in weniger als Sekunde zu berechnen.
Wir enthalten eine beschleunigte Umsetzung mit niedrigem Memory-Verbrauch.
Standardmäßig berechnen wir nur den Top -Singularwert (die spektrale Norm), aber dieser Code unterstützt das Berechnen von mehr SVs über das Argument --num_G_SVs .

Wichtige Unterschiede zwischen diesem Code und dem ursprünglichen Biggan

Wir verwenden die Optimierer-Einstellungen von SA-GAN (g_lr = 1e-4, d_lr = 4e-4, num_d_steps = 1 im Gegensatz zu Biggans g_lr = 5e-5, d_lr = 2e-4, num_d_steps = 2). Dies war die erste Ecke, die wir geschnitten haben, um die Trainingszeiten zu fällen.
Standardmäßig verwenden wir Cross-Replica-Batchnorm (auch bekannt als synchronisierte Batchnorm). Die beiden Varianten, die wir ausprobiert haben (eine benutzerdefinierte, naive und die in diesem Repo enthaltene), haben leicht unterschiedliche Gradienten (wenn auch identische Vorwärtspässe) als ein eingebauter Batchnorm, die ausreichend zu sein scheinen, um das Training zu lähmen.
Die Gradientenakkumulation bedeutet, dass wir die SV -Schätzungen und die BN -Statistik 8 -mal häufiger aktualisieren. Dies bedeutet, dass die BN -Statistiken den ständigen Statistiken viel näher sind und dass die Einfachwertschätzungen in der Regel genauer sind. Aus diesem Grund messen wir standardmäßig Metriken mit G im Testmodus (unter Verwendung des BatchNorm -Ausgangs -Statistikschätzungen anstatt stehende Statistiken wie im Papier zu berechnen). Wir unterstützen immer noch ständige Statistiken (siehe die sample.sh -Skripte). Dies könnte auch denkbar dazu führen, dass Gradienten aus den früheren Ansammlungen abgestanden sind, aber in der Praxis scheint dies kein Problem zu sein.
Die derzeit bereitgestellten Modelle wurden nicht mit orthogonaler Regularisierung geschult. Das Training ohne ortho reg scheint die Wahrscheinlichkeit zu erhöhen, dass Modelle nicht für die Kürzung zugänglich sind, aber es sieht so aus, als hätte dieses bestimmte Modell ein Gewinnerticket. Unabhängig davon bieten wir zwei hoch optimierte (schnelle und minimale Speicherverbrauch) ortho Reg -Implementierungen, die die Ortho Reg direkt berechnen. Gradienten.

Eine Notiz zum Design dieses Repo

Dieser Code ist von Grund auf so konzipiert, dass er als erweiterbare, hackbare Basis für den weiteren Forschungscode dient. Wir haben viel darüber nachgedacht, dass die Abstraktionen die richtige Dicke für die Forschung sind-nicht so dick, dass es undurchdringlich, aber nicht so dünn ist, dass sie nutzlos sind. Die Hauptidee ist, dass Sie, wenn Sie mit einem SOTA-Setup experimentieren und eine gewisse Änderung vornehmen möchten (probieren Sie Ihre eigene neue Verlustfunktion, Architektur, Selbstbekämpfung usw. aus), indem Sie dies einfach tun können, indem Sie Ihren Code an ein oder zwei Stellen fallen lassen, ohne sich um den Rest der Codebasis Sorgen zu machen. Dinge wie die Verwendung von Self.Which_conv und Functools.Partial in der Biggan.Py -Modelldefinition wurden berücksichtigt, ebenso wie das Design der Vererbung der Spectral Norm -Klasse.

Trotzdem ist dies eine etwas große Codebasis für ein einzelnes Projekt. Während wir versucht haben, mit den Kommentaren gründlich zu sein, können Sie bitte ein Problem oder eine Pull -Anfrage ansprechen, wenn es etwas gibt, das Sie für klarer, besser geschrieben oder besser umgestaltet haben könnten.

Feature -Anfragen

Möchten Sie an diesem Code arbeiten oder verbessern? Es gibt ein paar Dinge, von denen dieses Repo profitieren würde, aber die noch nicht funktionieren.

Synchronisierter Batchnorm (auch bekannt als Cross-Replica-Batchnorm). Wir haben zwei Varianten davon ausprobiert, aber aus irgendeinem unbekannten Grund hat es jedes Mal das Training verkrüppelt. Wir haben den Apex Syncbn nicht ausprobiert, da die Server meiner Schule auf alten Nvidia-Treibern sind, die ihn nicht unterstützen-APEX wäre wahrscheinlich ein guter Ausgangspunkt.
Gemischtes Präzisionstraining und Nutzung von Tensorkernen. Dieses Repo enthält eine naive ADAM-Implementierung mit gemischter Präzision, die früh im Training funktioniert, aber zu einem frühen Zusammenbruch führt und nichts unternimmt, um Tensorkerne zu aktivieren (es reduziert nur den Speicherverbrauch). Die Integration von Apex in diesen Code und die Verwendung seiner Schulungstechniken mit gemischten Vorräten, um Tensorkerne zu nutzen und den Speicherverbrauch zu verringern, kann zu erheblichen Geschwindigkeitsergebnissen führen.

MISSE Notizen

Siehe dieses Verzeichnis für ImageNet -Labels.

Wenn Sie diesen Code verwenden, zitieren Sie bitte

 @inproceedings{
brock2018large,
title={Large Scale {GAN} Training for High Fidelity Natural Image Synthesis},
author={Andrew Brock and Jeff Donahue and Karen Simonyan},
booktitle={International Conference on Learning Representations},
year={2019},
url={https://openreview.net/forum?id=B1xsqj09Fm},
}