StyleGAN1のから始める実装

チュートリアル

導入

この記事では、現在最も優れたGANの一つであるStyleGANについて説明します。論文「A Style-Based Generator Architecture for Generative Adversarial Networks」に基づいて、PyTorchを使用してクリーンでシンプルで読みやすい実装を行い、元の論文を可能な限り再現します。したがって、論文を読んだ場合、この実装は非常に似ているはずです。

このブログで使用するデータセットは、Kaggleから取得したデータセットで、256*192の解像度を持つ女性の上着16240点が含まれています。

前提知識

PyTorchを使用してStyleGANを扱う前に、以下の前提知識を確認してください：

  • ディープラーニングの基本知識
    コンボリューショナルニューラルネットワーク（CNN）の理解。
    生成対向ネットワーク（GANs）の熟悉、生成器、判別器、対向損失などの概念。

  • ハードウェア要件
    高速なトレーニングと推論のための強力なGPU（NVIDIA推奨）。
    GPU加速のためのCUDAツールキットがインストールされていること（cudaおよびcudnn）。

  • StyleGANに精通していること
    アーキテクチャの改善と主要な概念を理解するために、オリジナルのStyleGANまたはStyleGAN2の論文を読んでおくと役立つ。

必要なすべての依存関係を読み込む

まず、PyTorchを使用するためtorchをインポートし、nnをインポートします。これにより、ネットワークの作成と訓練が可能になり、最適化アルゴリズム（例：SGD、Adamなど）を実装するためのoptimパッケージをインポートできます。torchvisionからは、データの準備と変換を適用するためのdatasetsとtransformsをインポートします。

torch.nnからfunctionalをFとしてインポートし、画像をインターポレートしてアップサンプリングするため、DataLoaderをtorch.utils.dataからインポートし、いくつかのfakeサンプルを保存するためのsave_imageをtorchvision.utilsからインポートし、2のべき乗の逆表現を implementするためにmathのlog2をインポートし、線形代数のためのNumPyをインポートし、オペレーティングシステムとのインタラクションのためのosをインポートし、プログレスバーを表示するためのtqdmをインポートし、最後に結果を表示し、実際のものと比較するためのmatplotlib.pyplotをインポートします。

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

ハイパーパラメータ

  • 実際の画像のパスでDATASETを初期化します。
  • 画像サイズ8×8でトレーニングを開始指定します。
  • Cudaが利用可能な場合はデバイスを初期化し、CPUの場合はそれに従い、学習率を0.001に設定します。
  • 生成したい画像の解像度に応じてバッチサイズが異なるため、BATCH_SIZESを数値のリストで初期化し、VRAMに応じて変更できます。
  • 画像サイズを128に初期化し、CHANNELS_IMGを3に初期化します。なぜなら、128×128のRGB画像を生成するからです。
  • 元の論文では、Z_DIM、W_DIM、IN_CHANNELSを512として初期化していますが、私はVRAMの使用量を減らし、トレーニングを高速化するために256として初期化します。それを倍にすると、より良い結果が得られるかもしれません。
  • StyleGANでは、私たちが望むGANの損失関数を使用することができますので、私は論文Improved Training of Wasserstein GANsのWGAN-GPを使用しています。この損失にはλという名前のパラメータがあり、λ = 10に設定することが一般的です。
  • 各画像サイズごとにPROGRESSIVE_EPOCHSを30として初期化します。
DATASET                 = "Women clothes"
START_TRAIN_AT_IMG_SIZE = 8 #著者は4x4の画像ではなく8x8の画像から始めています
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)

データローダーを取得する

では、get_loaderという関数を作成して：

  • 画像にいくつかの変換を適用します（画像を希望する解像度にリサイズし、テンソルに変換し、いくつかの増加を適用し、最後にすべてのピクセルを-1から1の範囲に正規化します）。
  • 現在のバッチサイズをリストBATCH_SIZESを使って特定し、画像サイズ/4の2の逆数の整数指数をインデックスとして取り、これが実際にアダプティブミニバッチサイズを出力解像度に基づいて実装する方法です。
  • ImageFolderを使ってデータセットを準備します。なぜなら、それはすでにきれいに構造化されているからです。
  • データローダーを使用してミニバッチサイズを作成し、データセットとバッチサイズを指定し、データをシャッフルします。
  • 最後に、ローダーとデータセットを返します。
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

モデルの実装

さて、StyleGAN1のジェネレータとディスクリミネータ（ProGANとStyleGAN1はディスクリミネータのアーキテクチャが同じです）を実装し、論文の主要な特性を取り入れてみましょう。実装をコンパクトにし、同時に読みやすく理解しやすいようにします。特に以下のポイント：

  • ノイズマッピングネットワーク
  • アダプティブインスタンス正規化（AdaIN）
  • 段階的な成長

このチュートリアルでは、StyleGAN1を使用して画像を生成するだけで、スタイルミックスやスト Dichャスティック変動を実装はしませんが、それほど難しいことはありません。

各画像解像度におけるチャネル数を得るためにIN_CHANNELSに乗算する数を含む変数factorsを定義しましょう。

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

ノイズマッピングネットワーク

ノイズマッピングネットワークはZを取り入れ、8つの完全接続レイヤーに分けてアクティベーションを挟みます。そして、ProGAN（同じ研究者によって記述されたProGANとStyleGan）の作者が行うように、学習率を等価にしてください。

まず、WSLinear（加重スケールリニア）という名前のクラスをnn.Moduleから継承して構築しましょう。

  • init部分ではin_featuresとout_channelsを受け取り、リニアレイヤーを作成します。その後、スケールを定義し、それは2の平方根をin_featuresで割った値になります。現在のコラムレイヤーのバイアスを変数にコピーし、リニアレイヤーのバイアスをスケールしないためにその後バイアスを削除し、最後にリニアレイヤーを初期化します。
  • forward部分ではxを受け取り、行うこと全心がxにスケールを掛け、リシェイプしたバイアスを加えることです。
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

        # リニアレイヤーを初期化
        nn.init.normal_(self.linear.weight)
        nn.init.zeros_(self.bias)

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

次にMappingNetworkクラスを作成しましょう。

  • init部分ではz_dimとw_dinを受け取り、ネットワークマッピングを定義します。まずz_dimを正規化し、その後WSLinearとReLUをアクティベーション関数として8回繰り返します。
  • forward部分ではネットワークマッピングを返します。

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)

アダプティブインスタンス正規化（AdaIN）

さて、AdaINクラスを作成しましょう。

  • initの部分では、チャンネル、w_dimを送信し、インスタンス正規化の部分となるinstance_normを初期化し、スタイルスケールとスタイルバイアスをWSLinearでノイズマッピングネットワークWをチャンネルにマッピングするアダプティブな部分として初期化します。
  • forwardパスでは、xを送信し、インスタンス正規化を適用し、style_sclate * x + style_biasを返します。

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

ノイズ注入

さて、ノイズをジェネレーターに注入するためのクラスInjectNoiseを作成しましょう。

  • initの部分では、チャンネルを送信し、ランダムな正規分布からの重みを初期化し、nn.Parameterを使用してこれらの重みを最適化できるようにします。
  • forwardの部分では、画像xを送信し、ランダムなノイズを追加して返します。
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

役立つクラス

著者はStyleGANをKarras et alのProGANの公式実装之上に構築し、同じディスクリミネーターのアーキテクチャ、アダプティブミニバッチサイズ、ハイパーパラメータなどを使用しています。したがって、ProGANの実装から変更されていないクラスが多くあります。

この節では、ProGANのアーキテクチャから変更されていないクラスを作成します。

コードスニペット below に、クラス WSConv2d（加重スケーリングコンボリューションレイヤー）を Equalized Learning Rate ための conv レイヤー見つけることができます。

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

        convレイヤーを初期化
        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)

コードスニペット below に、ノイズマッピングネットワークの前に Z を正規化するためのクラス PixelNorm を見つけることができます。

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)

コードスニペット below に、ディスクリミネーターの作成を助けるクラス ConvBock を見つけることができます。

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

コードスニペット below に、ProGAN と同じディスクリミネータークラス Discriminatowich を見つけることができます。

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)

        ここでは、判別器のディスクリミネータからバックワイズに要因を戻すbecause the discriminator
        generatorから反転させることが必要だから、最初のprog_blockと
        rgbレイヤーを追加すると、入力サイズ1024x1024、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)
            )

        「initial_rgb」という少し混乱する名前ですが、これは4x4の入力サイズのためのRGBレイヤーです
        generatorのinitial_rgbを「ミラー」するためにこれを行いました
        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
        )  アベレージプールを使用したダウンサンプリング

        これは4x4の入力サイズのためのブロックです
        self.final_block = nn.Sequential(
            MiniBatch stdから結合するためにin_channelsに+1します
            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
            ),  線形レイヤーの代わりにこれを使用します
        )

    def fade_in(self, alpha, downscaled, out):
        """Used to fade in downscaled using avg pooling and output from CNN"""
        alphaは[0, 1]のスカラー値で、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])
        )
        各サンプル（すべてのチャネルとピクセルを跨いで）のstdを取得し、それを繰り返して
        シングルチャネルと結合し、画像に情報を追加します。この方法で判別器は
        バッチ/画像の変動について情報を得ます
        return torch.cat([x, batch_statistics], dim=1)

    def forward(self, x, alpha, steps):
        prog_blocksのリストをどこから始めるべきか、少し混乱するかもしれませんが
        最後のは4x4のためです。例えばsteps=1の場合、最後から二番目のものから始めるべきです
        input_sizeが8x8になるからです。steps==0の場合は最終ブロックを使用します
        
        cur_step = len(self.prog_blocks) - steps

        初期ステップとしてrgbから変換します。これは画像サイズに依存します
        それぞれにrgbレイヤーがあります
        out = self.leaky(self.rgb_layers[cur_step](x))

        if steps == 0:  例えば、画像が4x4の場合
            out = self.minibatch_std(out)
            return self.final_block(out).view(out.shape[0], -1)

        prog_blocksがチャネルを変更する可能性があるため、ダウンスケールでは前の/小さいサイズのrgb_layerを使用します
        私たちの場合、インデックスの+1に対応しています
        downscaled = self.leaky(self.rgb_layers[cur_step + 1](self.avg_pool(x)))
        out = self.avg_pool(self.prog_blocks[cur_step](out))

        fade_inはまずダウンスケールされたものと入力間で行います
        これは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)

ジェネレータ

ジェネレータアーキテクチャでは、いくつかのパターンが繰り返されるため、まずそれらのクラスを作成してコードをできるだけクリーンにします。クラス名をGenBlockとし、nn.Moduleから継承します。

  • init 部分では、in_channels、out_channels、w_dimを引数として受け取り、WSConv2dを使用してin_channelsをout_channelsにマッピングするconv1を初期化し、WSConv2dを使用してout_channelsをout_channelsにマッピングするconv2を初期化し、論文で使用されている0.2のスロープを持つLeaky ReLUをleakyとして初期化し、InjectNoiseを使用してinject_noise1、inject_noise2を初期化し、AdaINを使用してadain1、adain2を初期化します。
  • forward 部分では、xを引数として受け取り、conv1に渡してからinject_noise1とleakyに渡し、adain1で正規化し、その後conv2に渡してからinject_noise2とleakyに渡し、adain2で正規化します。最後に、xを返します。
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

これで、ジェネレータを作成するために必要なすべてのものが揃いました。


  • في الجزء init دعونا ن�始化 ‘starting_constant’ باستخدام متجه ثابت 4×4 (x 512 قناة في ورقة الأصل، و 256 في حالتنا) والذي يمر بiteración من الجينيراتور، يُرسم بواسطة ‘MappingNetwork’， initial_adain1، initial_adain2 بواسطة AdaIN، initial_noise1، initial_noise2 بواسطة InjectNoise، initial_conv بواسطة طبقة التحويل التي ترسم في_channels إلى نفسها، leaky بواسطة Leaky ReLU مع انحدار 0.2، initial_rgb بواسطة WSConv2d التي ترسم في_channels إلى img_channels والتي تكون 3 للRGB، prog_blocks بواسطة ModuleList() والذي سيحتوي على جميع الكتل التدرجية (نindicate مدخل/مخرج القنوات التيرجولية بضرب في_channels وهو 512 في الورقة و 256 في حالتنا في العوامل)، و rgb_blocks بواسطة ModuleList() والذي سيحتوي على جميع الكتل RGB.
  • لإدراج طبقات جديدة (جزء الأصل في ProGAN)، نضيف الجزء fade_in، الذي نرسل إليه alpha، scaled، و generated، ونقوم بإرجاع [tanh(alpha∗generated+(1−alpha)∗upscale)]، السبب في استخدامنا لتانх هو أن سيكون هو المخرج (الصورة المولدة) ونريد أن تكون البكسلات في نطاق بين 1 و -1.
  • 前方の部分では、ノイズ（Z_dim）、トレーニング中にゆっくりとフェードインするアルファ値（アルファは0から1の間）、現在取り扱っている解像度のステップ数を送信し、xをマップにパスして中間ノイズベクトルWを取得し、starting_constantをinitial_noise1に適用し、Wにinitial_adain1を適用します。次にinitial_convにパスし、再びleakyアクティベーション関数を使用してinitial_noise2を追加し、それとWにinitial_adain2を適用します。ステップが0かどうかを確認し、0の場合は初期RGBを通過させて終了します。それ以外の場合は、ステップ数をループし、各ループでアップスケーリング（upscaled）を行い、対応する解像度のプログレッシブブロック（out）を通過させます。最後に、フェードイン関数を返します。この関数はアルファ、final_out、RGBにマッピングしたfinal_upscaledを受け取ります。
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):  #インデックスエラーを防ぐためにfactors[i+1]の-1
            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):
        #アルファはスカラーで[0, 1]の範囲内にあり、 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)

        # upscaleのチャンネル数は変わらないが、
        # prog_blocksを通過したoutは変わる可能性がある。両方をRGBに変換できるように、
        # 異なるrgb_layersを使用します。
        # (steps-1)とstepsをそれぞれupscaled、outに対して使用します。
        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

以下のコードスニペットに、ジェネレータ gen、現在の解像度を特定するためのステップ数、および n=100 の数字を取る generate_examples 関数が含まれています。この関数の目的は、n個の偽画像を生成し、結果として保存することです。

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()

以下のコードスニペットに、WGAN-GP損失のための gradient_penalty 関数が含まれています。

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)

    # クリティックスコアの計算
    mixed_scores = critic(interpolated_images, alpha, train_step)
 
    # スコアの画像に対する勾配を取る
    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

トレーニング関数

トレーニング関数では、クリティック（ディスクリミネータ）、ジェネレータ gen、ローダー、データセット、ステップ、アルファ、およびジェネレータとクリティックのためのオプティマイザを送信します。

まず、DataLoaderで作成したすべてのミニバッチサイズをループし、画像だけを取り出します。なぜなら、ラベルは必要ないからです。

その後、ディスクリミネータ/Cリティックのトレーニングを設定し、E(critic(real)) – E(critic(fake))を最大化したいときに使用します。この方程式は、クリティックが実際の画像と偽画像をどれだけ区別できるかを意味します。

その後、生成器のトレーニングを設定し、E(critic(fake))。

最終的に、ループを更新し、fade_inのアルファ値を0から1の間に保ち、それを返します。

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()

        # アルファ値を更新し、1未満に保つ
        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

トレーニング

今、すべてが揃っているので、それらを合わせてStyleGANをトレーニングしましょう。

まず、生成器、ディスクリミネーター/クリティック、オプティマイザを初期化し、生成器とクリティックをトレーニングモードにします。その後、PROGRESSIVE_EPOCHSをループし、各ループでトレーニング関数をエポック数回呼び出し、いくつかのフェイク画像を生成し保存し、generate_examples関数を使って結果として保存し、最終的に次の画像解像度に進みます。

gen = Generator(
        Z_DIM, W_DIM, IN_CHANNELS, img_channels=CHANNELS_IMG
    ).to(DEVICE)
critic = Discriminator(IN_CHANNELS, img_channels=CHANNELS_IMG).to(DEVICE)
# オプティマイザを初期化
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()

# 設定ファイルで設定した画像サイズに対応するステップから始める
step = int(log2(START_TRAIN_AT_IMG_SIZE / 4))
for num_epochs in PROGRESSIVE_EPOCHS[step:]:
    alpha = 1e-5   # 非常に低いアルファ値から始める
    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  # 次の画像サイズに進む

結果

be

hopes that you can follow all the steps and gain a good understanding of how to correctly implement StyleGAN. Let’s now look at the results we get after training this model on this dataset with a resolution of 128×128.

Conclusion

In this article, we have created a clean, simple, and readable implementation of StyleGAN1 from scratch using PyTorch. We have replicated the original paper as closely as possible, so if you read the paper, the implementation should be almost identical.

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