#

MNISTとのWGAN-GP実験

10import torch
11
12from labml import experiment, tracker

#

ワッサーシュタインの実験から構成をインポートする

14from labml_nn.gan.wasserstein.experiment import Configs as OriginalConfigs

#

16from labml_nn.gan.wasserstein.gradient_penalty import GradientPenalty

#

設定クラス

オリジナルのGAN実装を拡張し、ディスクリミネーター（クリティック）損失計算をオーバーライドして勾配ペナルティを含めました。

19class Configs(OriginalConfigs):

#

グラデーションペナルティ係数 $λ$

28    gradient_penalty_coefficient: float = 10.0

#

30    gradient_penalty = GradientPenalty()

#

これは元のディスクリミネーター損失計算よりも優先され、勾配ペナルティも含まれます。

32    def calc_discriminator_loss(self, data: torch.Tensor):

#

$x$ グラデーションペナルティを計算するにはグラデーションをオンにする必要があります

38        data.requires_grad_()

#

[サンプル] $z \sim p (z)$

40        latent = self.sample_z(data.shape[0])

#

$D (x)$

42        f_real = self.discriminator(data)

#

$D (G_{θ} (z))$

44        f_fake = self.discriminator(self.generator(latent).detach())

#

ディスクリミネーターロスを取得

46        loss_true, loss_false = self.discriminator_loss(f_real, f_fake)

#

トレーニングモードでのグラデーションペナルティの計算

48        if self.mode.is_train:
49            gradient_penalty = self.gradient_penalty(data, f_real)
50            tracker.add("loss.gp.", gradient_penalty)
51            loss = loss_true + loss_false + self.gradient_penalty_coefficient * gradient_penalty

#

それ以外の場合はグラデーションペナルティをスキップ

53        else:
54            loss = loss_true + loss_false

#

ログのもの

57        tracker.add("loss.discriminator.true.", loss_true)
58        tracker.add("loss.discriminator.false.", loss_false)
59        tracker.add("loss.discriminator.", loss)
60
61        return loss

#

64def main():

#

コンフィグオブジェクトの作成

66    conf = Configs()

#

実験を作成

68    experiment.create(name='mnist_wassertein_gp_dcgan')

#

オーバーライド設定

70    experiment.configs(conf,
71                       {
72                           'discriminator': 'cnn',
73                           'generator': 'cnn',
74                           'label_smoothing': 0.01,
75                           'generator_loss': 'wasserstein',
76                           'discriminator_loss': 'wasserstein',
77                           'discriminator_k': 5,
78                       })

#

実験を開始してトレーニングループを実行する

81    with experiment.start():
82        conf.run()
83
84
85if __name__ == '__main__':
86    main()