Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Refactor basic GAN tutorial #357

Merged
merged 4 commits into from
Aug 1, 2024
Merged

Refactor basic GAN tutorial #357

Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
71d4fe8
fix logging issue
jhauret Aug 1, 2024
422204d
remove repeated calls to generator
jhauret Aug 1, 2024
ce02a32
more image samples
jhauret Aug 1, 2024
c39a740
a bit better
jhauret Aug 1, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
15 changes: 9 additions & 6 deletions lightning_examples/basic-gan/gan.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -88,7 +88,7 @@ def block(in_feat, out_feat, normalize=True):
layers = [nn.Linear(in_feat, out_feat)]
if normalize:
layers.append(nn.BatchNorm1d(out_feat, 0.8))
layers.append(nn.LeakyReLU(0.2, inplace=True))
layers.append(nn.LeakyReLU(0.01, inplace=True))
return layers

self.model = nn.Sequential(
Expand Down Expand Up @@ -193,15 +193,15 @@ def training_step(self, batch):
# log sampled images
sample_imgs = self.generated_imgs[:6]
grid = torchvision.utils.make_grid(sample_imgs)
self.logger.experiment.add_image("generated_images", grid, 0)
self.logger.experiment.add_image("train/generated_images", grid, self.current_epoch)

# ground truth result (ie: all fake)
# put on GPU because we created this tensor inside training_loop
valid = torch.ones(imgs.size(0), 1)
valid = valid.type_as(imgs)

# adversarial loss is binary cross-entropy
g_loss = self.adversarial_loss(self.discriminator(self(z)), valid)
g_loss = self.adversarial_loss(self.discriminator(self.generated_imgs), valid)
self.log("g_loss", g_loss, prog_bar=True)
self.manual_backward(g_loss)
optimizer_g.step()
Expand All @@ -222,7 +222,7 @@ def training_step(self, batch):
fake = torch.zeros(imgs.size(0), 1)
fake = fake.type_as(imgs)

fake_loss = self.adversarial_loss(self.discriminator(self(z).detach()), fake)
fake_loss = self.adversarial_loss(self.discriminator(self.generated_imgs.detach()), fake)

# discriminator loss is the average of these
d_loss = (real_loss + fake_loss) / 2
Expand All @@ -232,6 +232,9 @@ def training_step(self, batch):
optimizer_d.zero_grad()
self.untoggle_optimizer(optimizer_d)

def validation_step(self, batch, batch_idx):
pass

def configure_optimizers(self):
lr = self.hparams.lr
b1 = self.hparams.b1
Expand All @@ -247,7 +250,7 @@ def on_validation_epoch_end(self):
# log sampled images
sample_imgs = self(z)
grid = torchvision.utils.make_grid(sample_imgs)
self.logger.experiment.add_image("generated_images", grid, self.current_epoch)
self.logger.experiment.add_image("validation/generated_images", grid, self.current_epoch)


# %%
Expand All @@ -263,4 +266,4 @@ def on_validation_epoch_end(self):
# %%
# Start tensorboard.
# %load_ext tensorboard
# %tensorboard --logdir lightning_logs/
# %tensorboard --logdir lightning_logs/ --samples_per_plugin=images=60
Loading