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

19class Configs(OriginalConfigs):

28    gradient_penalty_coefficient: float = 10.0

30    gradient_penalty = GradientPenalty()

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

38        data.requires_grad_()

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

42        f_real = self.discriminator(data)

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