#

WGAN experiment with MNIST

9from labml import experiment
10
11from labml.configs import calculate

#

Import configurations from DCGAN experiment

13from labml_nn.gan.dcgan import Configs

#

Import Wasserstein GAN losses

16from labml_nn.gan.wasserstein import GeneratorLoss, DiscriminatorLoss

#

Set configurations options for Wasserstein GAN losses

19calculate(Configs.generator_loss, 'wasserstein', lambda c: GeneratorLoss())
20calculate(Configs.discriminator_loss, 'wasserstein', lambda c: DiscriminatorLoss())

#

23def main():

#

Create configs object

25    conf = Configs()

#

Create experiment

27    experiment.create(name='mnist_wassertein_dcgan', comment='test')

#

Override configurations

29    experiment.configs(conf,
30                       {
31                           'discriminator': 'cnn',
32                           'generator': 'cnn',
33                           'label_smoothing': 0.01,
34                           'generator_loss': 'wasserstein',
35                           'discriminator_loss': 'wasserstein',
36                       })

#

Start the experiment and run training loop

39    with experiment.start():
40        conf.run()
41
42
43if __name__ == '__main__':
44    main()