WGAN 使用 MNIST 进行实验

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

DCGAN 实验导入配置

13from labml_nn.gan.dcgan import Configs
16from labml_nn.gan.wasserstein import GeneratorLoss, DiscriminatorLoss

为 Wasserstein GAN 损耗设置配置选项

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

创建配置对象

25    conf = Configs()

创建实验

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

覆盖配置

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

开始实验并运行训练循环

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