wassname/pytorch-lightning
1
0
Fork 0
mirror of https://github.com/wassname/pytorch-lightning.git synced 2026-09-09 11:32:07 +08:00
Code Issues Packages Projects Releases Wiki Activity

fixed gan template (#528)

Browse Source 
* fixed gan template

* Update gan.py
This commit is contained in:
William Falcon
2019-12-04 08:28:46 -05:00
committed by GitHub
parent 218f0a5b4a
commit 6ba30a113d
1 changed files with 9 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
pl_examples/domain_templates/gan.py
+9
View File
@@ -112,7 +112,10 @@ class GAN(pl.LightningModule):
# self.logger.experiment.add_image('generated_images', grid, 0)
# ground truth result (ie: all fake)
# put on GPU because we created this tensor inside training_loop
valid = torch.ones(imgs.size(0), 1)
if self.on_gpu:
valid = valid.cuda(imgs.device.index)
# adversarial loss is binary cross-entropy
g_loss = self.adversarial_loss(self.discriminator(self.generated_imgs), valid)
@@ -130,10 +133,16 @@ class GAN(pl.LightningModule):
# how well can it label as real?
valid = torch.ones(imgs.size(0), 1)
if self.on_gpu:
valid = valid.cuda(imgs.device.index)
real_loss = self.adversarial_loss(self.discriminator(imgs), valid)
# how well can it label as fake?
fake = torch.zeros(imgs.size(0), 1)
if self.on_gpu:
fake = fake.cuda(imgs.device.index)
fake_loss = self.adversarial_loss(
self.discriminator(self.generated_imgs.detach()), fake)