10importtorch1112fromlabmlimportexperiment,tracker

14fromlabml\_nn.gan.wasserstein.experimentimportConfigsasOriginalConfigs

16fromlabml\_nn.gan.wasserstein.gradient\_penaltyimportGradientPenalty

19classConfigs(OriginalConfigs):

28gradient\_penalty\_coefficient:float=10.0

30gradient\_penalty=GradientPenalty()

32defcalc\_discriminator\_loss(self,data:torch.Tensor):

38data.requires\_grad\_()

40latent=self.sample\_z(data.shape[0])

42f\_real=self.discriminator(data)

44f\_fake=self.discriminator(self.generator(latent).detach())

46loss\_true,loss\_false=self.discriminator\_loss(f\_real,f\_fake)

48ifself.mode.is\_train:49gradient\_penalty=self.gradient\_penalty(data,f\_real)50tracker.add("loss.gp.",gradient\_penalty)51loss=loss\_true+loss\_false+self.gradient\_penalty\_coefficient\*gradient\_penalty

53else:54loss=loss\_true+loss\_false

57tracker.add("loss.discriminator.true.",loss\_true)58tracker.add("loss.discriminator.false.",loss\_false)59tracker.add("loss.discriminator.",loss)6061returnloss

64defmain():

66conf=Configs()

68experiment.create(name='mnist\_wassertein\_gp\_dcgan')

70experiment.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})

81withexperiment.start():82conf.run()838485if\_\_name\_\_=='\_\_main\_\_':86main()