Path: blob/master/labml_nn/gan/wasserstein/experiment.py
4956 views
"""1---2title: WGAN experiment with MNIST3summary: This experiment generates MNIST images using convolutional neural network.4---56# WGAN experiment with MNIST7"""8from labml import experiment910from labml.configs import calculate11# Import configurations from [DCGAN experiment](../dcgan/index.html)12from labml_nn.gan.dcgan import Configs1314# Import [Wasserstein GAN losses](./index.html)15from labml_nn.gan.wasserstein import GeneratorLoss, DiscriminatorLoss1617# Set configurations options for Wasserstein GAN losses18calculate(Configs.generator_loss, 'wasserstein', lambda c: GeneratorLoss())19calculate(Configs.discriminator_loss, 'wasserstein', lambda c: DiscriminatorLoss())202122def main():23# Create configs object24conf = Configs()25# Create experiment26experiment.create(name='mnist_wassertein_dcgan', comment='test')27# Override configurations28experiment.configs(conf,29{30'discriminator': 'cnn',31'generator': 'cnn',32'label_smoothing': 0.01,33'generator_loss': 'wasserstein',34'discriminator_loss': 'wasserstein',35})3637# Start the experiment and run training loop38with experiment.start():39conf.run()404142if __name__ == '__main__':43main()444546