StyleGAN1 von Grund auf implementieren

Tutorials

Einführung

Dieser Artikel behandelt eine der besten GANs heute, StyleGAN aus dem Papier Une Style-Based Generator Architecture for Generative Adversarial Networks, wir werden eine saubere, einfache und lesbare Implementierung davon mit PyTorch erstellen und versuchen, das Originalpapier so nah wie möglich nachzubilden, sodass wenn Sie das Papier lesen, die Implementierung pretty much identical sein sollte.

Der Datensatz, den wir in diesem Blog verwenden werden, ist dieser Datensatz von Kaggle, der 16240 Oberteile für Frauen mit einer Auflösung von 256*192 enthält.

Voraussetzungen

Bevor Sie sich in die Arbeit mit StyleGAN unter Verwendung von PyTorch stürzen, stellen Sie sicher, dass Sie die folgenden Voraussetzungen erfüllen:

Grundkenntnisse im Deep Learning
Verständnis von konvolutionellen neuronalen Netzen (CNNs).
Vertrautheit mit Generative Adversarial Networks (GANs), einschließlich Konzepte wie Generator, Diskriminator und adversarische Verluste.
Hardware-Anforderungen
Eine leistungsstarke GPU (NVIDIA empfohlen) für schnelleres Training und Inferenz.
Installiertes CUDA-Toolkit zur GPU-Beschleunigung (cuda und cudnn).
Vertrautheit mit StyleGAN
Es ist hilfreich, die ursprünglichen StyleGAN oder StyleGAN2 Papiere gelesen zu haben, um die Architekturverbesserungen und Schlüsselkonzepte zu verstehen.

Laden Sie alle Abhängigkeiten, die wir brauchen

Wir werden zunächst torch importieren, da wir PyTorch verwenden werden, und davon nn importieren. Das wird uns helfen, Netzwerke zu erstellen und zu trainieren, und uns auch erlaubt, optim zu importieren, ein Paket, das verschiedene Optimierungs-Algorithmen implementiert (z.B. sgd, adam, …). Von torchvision importieren wir datasets und transforms, um die Daten vorzubereiten und einige Transformationen anzuwenden.

Wir werden functional als F von torch.nn importieren, um die Bilder mit interpolate hochzurechnen, DataLoader von torch.utils.data, um Mini-Batch-Größen zu erstellen, save_image von torchvision.utils, um einige Fake-Beispiele zu speichern, und log2 von math, weil wir die inverse Darstellung der Potenz von 2 benötigen, um die adaptive Mini-Batch-Größe abhängig der Ausgabelösung zu implementieren, NumPy für lineare Algebra, os für die Interaktion mit dem Betriebssystem, tqdm, um Fortschrittsbalken anzuzeigen, und schließlich matplotlib.pyplot, um die Ergebnisse anzuzeigen und sie mit den echten zu vergleichen.

import torch
from torch import nn, optim
from torchvision import datasets, transforms
import torch.nn.functional as F
from torch.utils.data import DataLoader
from torchvision.utils import save_image
from math import log2
import numpy as np
import os
from tqdm import tqdm
import matplotlib.pyplot as plt

Hyperparameter

Initialisiere das DATASET durch den Pfad der echten Bilder.
Gebe die Startgröße für das Training bei 8×8 Bildgröße an.
Initialisiere das Gerät mit Cuda, falls verfügbar, andernfalls mit CPU, und den Lernrate auf 0.001.
Die Batch-Größe wird je nach der Auflösung der zu generierenden Bilder unterschiedlich sein, daher initialisieren wir BATCH_SIZES mit einer Liste von Zahlen, die du abhängig von deinem VRAM ändern kannst.
Initialisiere image_size auf 128 und CHANNELS_IMG auf 3, weil wir 128×128 RGB-Bilder generieren werden.
In der ursprünglichen Arbeit initialisieren sie Z_DIM, W_DIM und IN_CHANNELS mit 512, aber ich initialisiere sie stattdessen mit 256, um weniger VRAM zu verwenden und das Training zu beschleunigen. Wir könnten vielleicht sogar bessere Ergebnisse erzielen, wenn wir sie verdoppeln.
Für StyleGAN können wir jede der GAN-Verlustfunktionen verwenden, die wir wollen,also verwende ich WGAN-GP aus dem Papier Improved Training of Wasserstein GANs. Diese Verlustfunktion enthält einen Parameter namens λ und es ist üblich, λ = 10 zu setzen.
Initialisiere PROGRESSIVE_EPOCHS mit 30 für jede Bildgröße.

DATASET                 = "Women clothes"
START_TRAIN_AT_IMG_SIZE = 8 #Die Autoren beginnen mit 8x8-Bildern anstatt mit 4x4
DEVICE                  = "cuda" if torch.cuda.is_available() else "cpu"
LEARNING_RATE           = 1e-3
BATCH_SIZES             = [256, 128, 64, 32, 16, 8]
CHANNELS_IMG            = 3
Z_DIM                   = 256
W_DIM                   = 256
IN_CHANNELS             = 256
LAMBDA_GP               = 10
PROGRESSIVE_EPOCHS      = [30] * len(BATCH_SIZES)

Datenlader abrufen

Jetzt erstellen wir eine Funktion get_loader, um:

Einige Transformationen auf die Bilder anzuwenden (die Bilder auf die gewünschte Auflösung verkleinern, sie in Tensoren umwandeln, dann einige Erweiterungen anwenden und schließlich normalisieren, sodass alle Pixel im Bereich von -1 bis 1 liegen).
Die aktuelle Batchgröße zu identifizieren, indem die Liste BATCH_SIZES verwendet wird, und als Index die ganzzahlige Darstellung der inverse Potenz von 2 der Bildgröße/4 zu nehmen. Und das ist tatsächlich, wie wir die adaptive Minibatch-Größe abhängig der Ausgabelösung implementieren.
Das Dataset vorzubereiten, indem ImageFolder verwendet wird, weil es bereits in einer schönen Weise strukturiert ist.
Erstellen Sie Mini-Batch-Größen mit DataLoader, der den Datensatz und die Batch-Größe übernimmt und die Daten mischt.
Schließlich geben Sie den Loader und den Datensatz zurück.

def get_loader(image_size):
    transform = transforms.Compose(
        [
            transforms.Resize((image_size, image_size)),
            transforms.ToTensor(),
            transforms.RandomHorizontalFlip(p=0.5),
            transforms.Normalize(
                [0.5 for _ in range(CHANNELS_IMG)],
                [0.5 for _ in range(CHANNELS_IMG)],
            ),
        ]
    )
    batch_size = BATCH_SIZES[int(log2(image_size / 4))]
    dataset = datasets.ImageFolder(root=DATASET, transform=transform)
    loader = DataLoader(
        dataset,
        batch_size=batch_size,
        shuffle=True,
    )
    return loader, dataset

Modellimplementierung

Jetzt implementieren wir den StyleGAN1-Generator und -Diskriminator (ProGAN und StyleGAN1 haben die gleiche Diskriminatorarchitektur) mit den Schlüsselattributen aus dem Papier. Wir werden versuchen, die Implementierung kompakt zu gestalten, aber auch lesbar und verständlich zu halten. Insbesondere die folgenden Schlüsselpunkte:

Rauschabbildungsnetzwerk
Adaptive Instanznormalisierung (AdaIN)
Progressives Wachsen

In diesem Tutorial werden wir nur Bilder mit StyleGAN1 generieren und nicht die Stil mieszene und stochastische Variation implementieren, aber das sollte nicht schwer sein.

Definieren wir eine Variable mit dem Namen factors, die die Zahlen enthält, die mit IN_CHANNELS multipliziert werden, um die gewünschte Anzahl von Kanälen in jeder Bildauflösung zu erhalten.

factors = [1, 1, 1, 1, 1 / 2, 1 / 4, 1 / 8, 1 / 16, 1 / 32]

Rauschabbildungsnetzwerk

Das RauschORKartierungsnetzwerk nimmt Z und führt es durch acht vollständig verbundene Schichten, die durch einige Aktivierung getrennt sind. Und vergessen Sie nicht, die Lernrate zu equalisieren, wie die Autoren in ProGAN (ProGAN und StyleGan von denselben Forschern verfasst) tun.

Lassen Sie uns zunächst eine Klasse mit dem Namen WSLinear (gewichtete skalierte Lineare) erstellen, die von nn.Module erben wird.

In der init-Teil geben wir in_features und out_channels weiter. Erstellen Sie eine lineare Ebene, dann definieren wir eine Skalierung, die gleich dem Quadratwurzel von 2 geteilt durch in_features ist, wir kopieren den Bias der aktuellen Spaltenschicht in eine Variable, weil wir den Bias der linearen Schicht nicht skalieren möchten, dann entfernen wir ihn, schließlich initialisieren wir die lineare Ebene.
In der forward-Teil senden wir x und alles, was wir tun werden, ist x mit der Skalierung zu multiplizieren und den Bias nach dem Umformen hinzuzufügen.

class WSLinear(nn.Module):
    def __init__(
        self, in_features, out_features,
    ):
        super(WSLinear, self).__init__()
        self.linear = nn.Linear(in_features, out_features)
        self.scale = (2 / in_features)**0.5
        self.bias = self.linear.bias
        self.linear.bias = None

        # initialisiere lineare Ebene
        nn.init.normal_(self.linear.weight)
        nn.init.zeros_(self.bias)

    def forward(self, x):
        return self.linear(x * self.scale) + self.bias

Jetzt erstellen wir die MappingNetwork-Klasse.

In der init-Teil senden wir z_dim und w_din, und wir definieren das Netzwerk-Mapping, das zuerst z_dim normalisiert, gefolgt von acht WSLInear und ReLU als Aktivierungsfunktionen.
In der forward-Teil geben wir das Netzwerk-Mapping zurück.

class MappingNetwork(nn.Module):
    def __init__(self, z_dim, w_dim):
        super().__init__()
        self.mapping = nn.Sequential(
            PixelNorm(),
            WSLinear(z_dim, w_dim),
            nn.ReLU(),
            WSLinear(w_dim, w_dim),
            nn.ReLU(),
            WSLinear(w_dim, w_dim),
            nn.ReLU(),
            WSLinear(w_dim, w_dim),
            nn.ReLU(),
            WSLinear(w_dim, w_dim),
            nn.ReLU(),
            WSLinear(w_dim, w_dim),
            nn.ReLU(),
            WSLinear(w_dim, w_dim),
            nn.ReLU(),
            WSLinear(w_dim, w_dim),
        )

    def forward(self, x):
        return self.mapping(x)

Anpassende Instanznormalisierung (AdaIN)

Jetzt erstellen wir die AdaIN-Klasse

In der INIT-Teil senden wir Kanäle, w_dim und initialisieren instance_norm, das der Instanznormalisierungsteil sein wird, und initialisieren style_scale und style_bias, die die adaptiven Teile mit WSLinear sein werden, die das Noise Mapping Network W in Kanäle abbilden.
In der FORWARD-Pass senden wir x, wenden die Instanznormalisierung darauf an und geben style_scale * x + style_bias zurück.

class AdaIN(nn.Module):
    def __init__(self, channels, w_dim):
        super().__init__()
        self.instance_norm = nn.InstanceNorm2d(channels)
        self.style_scale = WSLinear(w_dim, channels)
        self.style_bias = WSLinear(w_dim, channels)

    def forward(self, x, w):
        x = self.instance_norm(x)
        style_scale = self.style_scale(w).unsqueeze(2).unsqueeze(3)
        style_bias = self.style_bias(w).unsqueeze(2).unsqueeze(3)
        return style_scale * x + style_bias

Geräusch injizieren

Jetzt erstellen wir die Klasse InjectNoise, um das Geräusch in den Generator zu injizieren

In der INIT-Teil haben wir Kanäle gesendet und das Gewicht aus einer zufälligen Normalverteilung initialisiert und nn.Parameter verwendet, damit diese Gewichte optimiert werden können
In der FORWARD-Teil senden wir ein Bild x und geben es mit zufälligem Geräusch hinzugefügt zurück

class InjectNoise(nn.Module):
    def __init__(self, channels):
        super().__init__()
        self.weight = nn.Parameter(torch.zeros(1, channels, 1, 1))

    def forward(self, x):
        noise = torch.randn((x.shape[0], 1, x.shape[2], x.shape[3]), device=x.device)
        return x + self.weight * noise

Hilfreiche Klassen

Die Autoren haben StyleGAN auf der offiziellen Implementierung von ProGAN von Karras et al aufgebaut, sie verwenden die gleiche Diskriminatorarchitektur, adaptive Minibatch-Größe, Hyperparameter usw. Daher gibt es viele Klassen, die von der ProGAN-Implementierung gleich bleiben.

In diesem Abschnitt werden wir die Klassen erstellen, die von der ProGAN-Architektur nicht geändert werden.

Im folgenden Code-Schnipptel finden Sie die Klasse WSConv2d (gewichtete skalierte Konvolutionslage) für die Equalized Learning Rate für die Konvolutionslagen.

class WSConv2d(nn.Module):
    def __init__(
        self, in_channels, out_channels, kernel_size=3, stride=1, padding=1
    ):
        super(WSConv2d, self).__init__()
        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding)
        self.scale = (2 / (in_channels * (kernel_size ** 2))) ** 0.5
        self.bias = self.conv.bias
        self.conv.bias = None

        # initialisiere Konvolutionslage
        nn.init.normal_(self.conv.weight)
        nn.init.zeros_(self.bias)

    def forward(self, x):
        return self.conv(x * self.scale) + self.bias.view(1, self.bias.shape[0], 1, 1)

Im folgenden Code-Schnipptel finden Sie die Klasse PixelNorm zur Normalisierung von Z vor dem Noise Mapping Network.

class PixelNorm(nn.Module):
    def __init__(self):
        super(PixelNorm, self).__init__()
        self.epsilon = 1e-8

    def forward(self, x):
        return x / torch.sqrt(torch.mean(x ** 2, dim=1, keepdim=True) + self.epsilon)

Im folgenden Code-Schnipptel finden Sie die Klasse ConvBlock, die uns dabei hilft, den Diskriminator zu erstellen.

class ConvBlock(nn.Module):
    def __init__(self, in_channels, out_channels):
        super(ConvBlock, self).__init__()
        self.conv1 = WSConv2d(in_channels, out_channels)
        self.conv2 = WSConv2d(out_channels, out_channels)
        self.leaky = nn.LeakyReLU(0.2)

    def forward(self, x):
        x = self.leaky(self.conv1(x))
        x = self.leaky(self.conv2(x))
        return x

Im folgenden Code-Schnipptel finden Sie die Klasse Discriminatowich, die derselben wie in ProGAN ist.

class Discriminator(nn.Module):
    def __init__(self, in_channels, img_channels=3):
        super(Discriminator, self).__init__()
        self.prog_blocks, self.rgb_layers = nn.ModuleList([]), nn.ModuleList([])
        self.leaky = nn.LeakyReLU(0.2)

        hier arbeiten wir rückwärts von Faktoren, weil der Diskriminator
        von dem Generator gespiegelt werden sollte. Also der erste prog_block und
        die erste rgb Ebene, die wir hinzufügen, funktioniert für die Eingangsgröße 1024x1024, dann 512->256-> etc
        for i in range(len(factors) - 1, 0, -1):
            conv_in = int(in_channels * factors[i])
            conv_out = int(in_channels * factors[i - 1])
            self.prog_blocks.append(ConvBlock(conv_in, conv_out))
            self.rgb_layers.append(
                WSConv2d(img_channels, conv_in, kernel_size=1, stride=1, padding=0)
            )

        vielleicht verwirrendes Name "initial_rgb", dies ist nur die RGB Ebene für die 4x4 Eingangsgröße
        das haben wir gemacht, um den Generator initial_rgb zu spiegeln
        self.initial_rgb = WSConv2d(
            img_channels, in_channels, kernel_size=1, stride=1, padding=0
        )
        self.rgb_layers.append(self.initial_rgb)
        self.avg_pool = nn.AvgPool2d(
            kernel_size=2, stride=2
        )  Downsampling mit Durchschnittspooling

        das ist der Block für die 4x4 Eingangsgröße
        self.final_block = nn.Sequential(
            +1 zu in_channels, weil wir von MiniBatch std konkatenieren
            WSConv2d(in_channels + 1, in_channels, kernel_size=3, padding=1),
            nn.LeakyReLU(0.2),
            WSConv2d(in_channels, in_channels, kernel_size=4, padding=0, stride=1),
            nn.LeakyReLU(0.2),
            WSConv2d(
                in_channels, 1, kernel_size=1, padding=0, stride=1
            ),  wir verwenden das stattdessen von linearer Ebene
        )

    def fade_in(self, alpha, downscaled, out):
        """Used to fade in downscaled using avg pooling and output from CNN"""
        alpha sollte Skalar innerhalb [0, 1] sein, und upscale.shape == generated.shape
        return alpha * out + (1 - alpha) * downscaled

    def minibatch_std(self, x):
        batch_statistics = (
            torch.std(x, dim=0).mean().repeat(x.shape[0], 1, x.shape[2], x.shape[3])
        )
        wir nehmen die Std für jedes Beispiel (über alle Kanäle und Pixel) und wiederholen es
        für einen einzelnen Kanal und konkatenieren es mit dem Bild. Auf diese Weise erhält der Diskriminator
        Informationen über die Variation im Batch/Bild
        return torch.cat([x, batch_statistics], dim=1)

    def forward(self, x, alpha, steps):
        wo wir in der Liste der prog_blocks beginnen sollten, vielleicht ein bisschen verwirrend, aber
        das letzte ist für die 4x4. Also Beispiel, sagen wir steps=1, dann sollten wir
        bei dem vorletzten beginnen, weil input_size 8x8 sein wird. Wenn steps==0 verwenden wir einfach
        den letzten Block
        cur_step = len(self.prog_blocks) - steps

        Umwandlung von rgb als erster Schritt, das wird von
        der Bildgröße abhängen (jedes wird seine eigene rgb Ebene haben)
        out = self.leaky(self.rgb_layers[cur_step](x))

        if steps == 0:  nämlich, das Bild ist 4x4
            out = self.minibatch_std(out)
            return self.final_block(out).view(out.shape[0], -1)

        weil prog_blocks die Kanäle ändern könnte, verwenden wir für das Herunter skalieren die rgb_layer
        von der vorherigen/kleineren Größe, die in unserem Fall korrespondiert zu +1 in der Indizierung
        downscaled = self.leaky(self.rgb_layers[cur_step + 1](self.avg_pool(x)))
        out = self.avg_pool(self.prog_blocks[cur_step](out))

        der fade_in wird zuerst zwischen dem downscaled und dem Eingang durchgeführt
        das ist entgegengesetzt zum Generator
        out = self.fade_in(alpha, downscaled, out)

        for step in range(cur_step + 1, len(self.prog_blocks)):
            out = self.prog_blocks[step](out)
            out = self.avg_pool(out)

        out = self.minibatch_std(out)
        return self.final_block(out).view(out.shape[0], -1)

Generator

In der Generatorarchitektur haben wir einige Muster, die sich wiederholen, therefore erstellen wir zuerst eine Klasse dafür, um unseren Code so sauber wie möglich zu gestalten. Nennen wir die Klasse GenBlock, die von nn.Module erblickt.

In der init-Methode übergeben wir in_channels, out_channels und w_dim, initialisieren dann conv1 mit WSConv2d, der in_channels auf out_channels abbildet, conv2 mit WSConv2d, der out_channels auf out_channels abbildet, leaky mit Leaky ReLU mit einer Neigung von 0.2, wie sie im Papier verwenden, inject_noise1 und inject_noise2 mit InjectNoise, adain1 und adain2 mit AdaIN.
In der forward-Methode übergeben wir x und leiten es durch conv1, dann durch inject_noise1 mit leaky, normalisieren es mit adain1 und leiten es erneut durch conv2, dann durch inject_noise2 mit leaky und normalisieren es mit adain2. Und schließlich geben wir x zurück.

class GenBlock(nn.Module):
    def __init__(self, in_channels, out_channels, w_dim):
        super(GenBlock, self).__init__()
        self.conv1 = WSConv2d(in_channels, out_channels)
        self.conv2 = WSConv2d(out_channels, out_channels)
        self.leaky = nn.LeakyReLU(0.2, inplace=True)
        self.inject_noise1 = InjectNoise(out_channels)
        self.inject_noise2 = InjectNoise(out_channels)
        self.adain1 = AdaIN(out_channels, w_dim)
        self.adain2 = AdaIN(out_channels, w_dim)

    def forward(self, x, w):
        x = self.adain1(self.leaky(self.inject_noise1(self.conv1(x))), w)
        x = self.adain2(self.leaky(self.inject_noise2(self.conv2(x))), w)
        return x

Jetzt haben wir alles, was wir brauchen, um den Generator zu erstellen.

Im init-Teil initialisieren wir die ‘starting_constant’ durch einen Tensor mit der Konstanten 4 x 4 (x 512 Kanälen für denOriginalartikel und 256 in unserem Fall), der durch eine Iteration des Generators,abbildung durch ‘MappingNetwork’, initial_adain1, initial_adain2 durch AdaIN, initial_noise1, initial_noise2 durch InjectNoise, initial_conv durch eine Konvolutionslage, die in_channels auf sich selbst abbildet, leaky durch Leaky ReLU mit einer Neigung von 0.2, initial_rgb durch WSConv2d, die in_channels auf img_channelsabbildet, was für RGB 3 ist, prog_blocks durch ModuleList(), die alle progressiven Blöcke enthalten wird (wir geben die Kanäle der Konvolutionseingabe/Ausgabe durch Multiplikation mit in_channels an, das im Artikel 512 und in unserem Fall 256 ist, mit Faktoren), und rgb_blocks durch ModuleList(), die alle RGB-Blöcke enthalten wird.
Um neue Ebenen einblenden zu können (ein ursprüngliches Element von ProGAN), fügen wir den fade_in-Teil hinzu, dem wir alpha, scaled und generated übergeben, und wir geben [tanh(alpha∗generated+(1−alpha)∗upscale)] zurück. Der Grund, warum wir tanh verwenden, ist, dass dies die Ausgabe (das generierte Bild) sein wird und wir möchten, dass die Pixel im Bereich zwischen 1 und -1 liegen.
In der vorward-Teil senden wir das Rauschen (Z_dim), den Alpha-Wert, der während des Trainings langsam einblendet (Alpha liegt zwischen 0 und 1), und die Schritte, die die aktuelle Auflösung sind, mit der wir arbeiten. Wir übergeben x in die Karte, um den intermediären Rauschenvektor W zu erhalten, übergeben starting_constant an initial_noise1, wenden es an und initial_adain1 für W an, dann geben wir es in initial_conv weiter, und wieder.addieren wir initial_noise2 für es mit leaky als Aktivierungsfunktion und wenden es und W initial_adain2 an. Dann überprüfen wir, ob steps = 0 ist, wenn ja, dann möchten wir es nur durch das initial RGB laufen lassen und wir sind fertig, andernfalls wiederholen wir die Anzahl der Schritte und in jeder Schleife skalieren wir hoch (upscaled) und laufen durch den progressiven Block, der dieser Auflösung entspricht (out). Am Ende geben wir fade_in zurück, das alpha, final_out und final_upscaled nach der Abbildung auf RGB nimmt.

class Generator(nn.Module):
    def __init__(self, z_dim, w_dim, in_channels, img_channels=3):
        super(Generator, self).__init__()
        self.starting_constant = nn.Parameter(torch.ones((1, in_channels, 4, 4)))
        self.map = MappingNetwork(z_dim, w_dim)
        self.initial_adain1 = AdaIN(in_channels, w_dim)
        self.initial_adain2 = AdaIN(in_channels, w_dim)
        self.initial_noise1 = InjectNoise(in_channels)
        self.initial_noise2 = InjectNoise(in_channels)
        self.initial_conv = nn.Conv2d(in_channels, in_channels, kernel_size=3, stride=1, padding=1)
        self.leaky = nn.LeakyReLU(0.2, inplace=True)

        self.initial_rgb = WSConv2d(
            in_channels, img_channels, kernel_size=1, stride=1, padding=0
        )
        self.prog_blocks, self.rgb_layers = (
            nn.ModuleList([]),
            nn.ModuleList([self.initial_rgb]),
        )

        for i in range(len(factors) - 1):  # -1, um einen Indexfehler wegen factors[i+1] zu verhindern
            conv_in_c = int(in_channels * factors[i])
            conv_out_c = int(in_channels * factors[i + 1])
            self.prog_blocks.append(GenBlock(conv_in_c, conv_out_c, w_dim))
            self.rgb_layers.append(
                WSConv2d(conv_out_c, img_channels, kernel_size=1, stride=1, padding=0)
            )

    def fade_in(self, alpha, upscaled, generated):
        # Alpha sollte eine Skalare innerhalb von [0, 1] sein, und upscale.shape == generated.shape
        return torch.tanh(alpha * generated + (1 - alpha) * upscaled)

    def forward(self, noise, alpha, steps):
        w = self.map(noise)
        x = self.initial_adain1(self.initial_noise1(self.starting_constant), w)
        x = self.initial_conv(x)
        out = self.initial_adain2(self.leaky(self.initial_noise2(x)), w)

        if steps == 0:
            return self.initial_rgb(x)

        for step in range(steps):
            upscaled = F.interpolate(out, scale_factor=2, mode="bilinear")
            out = self.prog_blocks[step](upscaled, w)

        # Die Anzahl der Kanäle in upscale bleibt gleich, während
        # out, das durch prog_blocks bewegt wurde, sich ändern könnte. Um sicherzustellen
        # wir können beide in RGB umwandeln, verwenden wir不同的rgb_layers
        # (steps-1) und steps für upscaled, out respectively
        final_upscaled = self.rgb_layers[steps - 1](upscaled)
        final_out = self.rgb_layers[steps](out)
        return self.fade_in(alpha, final_upscaled, final_out)

Utils

In dem folgenden Code-Schnipsel finden Sie die Funktion generate_examples, die den Generator gen, die Anzahl der Schritte zur Identifizierung der aktuellen Auflösung und eine Zahl n=100 akzeptiert. Das Ziel dieser Funktion ist es, n gefälschte Bilder zu generieren und diese als Ergebnis zu speichern.

def generate_examples(gen, steps, n=100):

    gen.eval()
    alpha = 1.0
    for i in range(n):
        with torch.no_grad():
            noise = torch.randn(1, Z_DIM).to(DEVICE)
            img = gen(noise, alpha, steps)
            if not os.path.exists(f'saved_examples/step{steps}'):
                os.makedirs(f'saved_examples/step{steps}')
            save_image(img*0.5+0.5, f"saved_examples/step{steps}/img_{i}.png")
    gen.train()

In dem folgenden Code-Schnipsel finden Sie die Funktion gradient_penalty für die WGAN-GP-Verlustfunktion.

def gradient_penalty(critic, real, fake, alpha, train_step, device="cpu"):
    BATCH_SIZE, C, H, W = real.shape
    beta = torch.rand((BATCH_SIZE, 1, 1, 1)).repeat(1, C, H, W).to(device)
    interpolated_images = real * beta + fake.detach() * (1 - beta)
    interpolated_images.requires_grad_(True)

    # Berechne Kritiker-Werte
    mixed_scores = critic(interpolated_images, alpha, train_step)
 
    # Nehme den Gradienten der Werte gegenüber den Bildern
    gradient = torch.autograd.grad(
        inputs=interpolated_images,
        outputs=mixed_scores,
        grad_outputs=torch.ones_like(mixed_scores),
        create_graph=True,
        retain_graph=True,
    )[0]
    gradient = gradient.view(gradient.shape[0], -1)
    gradient_norm = gradient.norm(2, dim=1)
    gradient_penalty = torch.mean((gradient_norm - 1) ** 2)
    return gradient_penalty

Trainingsfunktion

Für die Trainingsfunktion übermitteln wir den Kritiker (der Diskriminator), den Generator gen, den Loader, den Datensatz, den Schritt, Alpha und den Optimierer für den Generator und den Kritiker.

Wir beginnen mit einer Schleife über alle Mini-Batch-Größen, die wir mit dem DataLoader erstellen, und nehmen nur die Bilder, da wir keine Label benötigen.

Dann richten wir das Training für den Diskriminator\Critic ein, wenn wir E(critic(real)) – E(critic(fake)) maximieren möchten. Diese Gleichung bedeutet, wie gut der Kritiker zwischen realen und gefälschten Bildern unterscheiden kann.

Nachdem wir das setup für den Generator durchgeführt haben, wenn wir E(critic(fake)).

maximieren möchten, aktualisieren wir die Schleife und den Alpha-Wert für fade_in und stellen sicher, dass er zwischen 0 und 1 liegt, und geben ihn zurück.

def train_fn(
    critic,
    gen,
    loader,
    dataset,
    step,
    alpha,
    opt_critic,
    opt_gen,
):
    loop = tqdm(loader, leave=True)

    for batch_idx, (real, _) in enumerate(loop):
        real = real.to(DEVICE)
        cur_batch_size = real.shape[0]


        noise = torch.randn(cur_batch_size, Z_DIM).to(DEVICE)

        fake = gen(noise, alpha, step)
        critic_real = critic(real, alpha, step)
        critic_fake = critic(fake.detach(), alpha, step)
        gp = gradient_penalty(critic, real, fake, alpha, step, device=DEVICE)
        loss_critic = (
            -(torch.mean(critic_real) - torch.mean(critic_fake))
            + LAMBDA_GP * gp
            + (0.001 * torch.mean(critic_real ** 2))
        )

        critic.zero_grad()
        loss_critic.backward()
        opt_critic.step()

        gen_fake = critic(fake, alpha, step)
        loss_gen = -torch.mean(gen_fake)

        gen.zero_grad()
        loss_gen.backward()
        opt_gen.step()

        # Alpha aktualisieren und sicherstellen, dass er kleiner als 1 ist
        alpha += cur_batch_size / (
            (PROGRESSIVE_EPOCHS[step] * 0.5) * len(dataset)
        )
        alpha = min(alpha, 1)

        loop.set_postfix(
            gp=gp.item(),
            loss_critic=loss_critic.item(),
        )


    return alpha

Training

Jetzt, da wir alles haben, lassen wir uns an die Arbeit machen und unseren StyleGAN trainieren.

Wir beginnen mit der Initialisierung des Generators, des Diskriminators/Kritikers und der Optimierer,然后将 Generator und Kritiker in den Trainingsmodus versetzen, dann iterieren wir über PROGRESSIVE_EPOCHS und rufen in jeder Schleife die Trainingsfunktion die Anzahl der Epochen auf, erstellen einige fake-Bilder und speichern sie mithilfe der generate_examples-Funktion und schreiten schließlich zur nächsten Bildauflösung vor.

gen = Generator(
        Z_DIM, W_DIM, IN_CHANNELS, img_channels=CHANNELS_IMG
    ).to(DEVICE)
critic = Discriminator(IN_CHANNELS, img_channels=CHANNELS_IMG).to(DEVICE)
# Optimierer initialisieren
opt_gen = optim.Adam([{"params": [param for name, param in gen.named_parameters() if "map" not in name]},
                        {"params": gen.map.parameters(), "lr": 1e-5}], lr=LEARNING_RATE, betas=(0.0, 0.99))
opt_critic = optim.Adam(
    critic.parameters(), lr=LEARNING_RATE, betas=(0.0, 0.99)
)


gen.train()
critic.train()

# bei dem Schritt starten, der der in der Konfiguration gesetzten Bildgröße entspricht
step = int(log2(START_TRAIN_AT_IMG_SIZE / 4))
for num_epochs in PROGRESSIVE_EPOCHS[step:]:
    alpha = 1e-5   # mit sehr niedrigem Alpha starten
    loader, dataset = get_loader(4 * 2 ** step)  
    print(f"Current image size: {4 * 2 ** step}")

    for epoch in range(num_epochs):
        print(f"Epoch [{epoch+1}/{num_epochs}]")
        alpha = train_fn(
            critic,
            gen,
            loader,
            dataset,
            step,
            alpha,
            opt_critic,
            opt_gen
        )

    generate_examples(gen, step)
    step += 1  # zur nächsten Bildgröße fortschreiten

Ergebnis

Hoffentlich können Sie alle Schritte folgen und eine gute Vorstellung davon bekommen, wie man StyleGAN auf die richtige Weise implementiert.Nun schauen wir uns die Ergebnisse an, die wir nach dem Training dieses Modells in diesem Datenbestand mit einer Auflösung von 128*x 128 erzielen.

Schlussfolgerung

In diesem Artikel haben wir eine saubere, einfache und lesbare Implementierung von StyleGAN1 von Grund auf mit PyTorch erstellt. Wir haben die Originalarbeit so nah wie möglich nachgeahmt, sodass die Implementierung, wenn Sie die Arbeit lesen, Pretty viel identisch sein sollte.

Source:
https://www.digitalocean.com/community/tutorials/implementation-stylegan-from-scratch