Spaces:
Sleeping
Sleeping
| import torch | |
| import torch.nn as nn | |
| import lightning as L | |
| import torch.optim as optim | |
| from models.generator import Generator | |
| from models.discriminator import Discriminator | |
| from utility.helper import save_some_examples | |
| class Pix2Pix(L.LightningModule): | |
| def __init__(self, in_channels, learning_rate, l1_lambda, features_generator, features_discriminator, display_step): | |
| super().__init__() | |
| self.automatic_optimization = False | |
| self.gen = Generator( | |
| in_channels=in_channels, | |
| features=features_generator | |
| ) | |
| self.disc = Discriminator( | |
| in_channels=in_channels, | |
| features=features_discriminator | |
| ) | |
| self.loss_fn = nn.BCEWithLogitsLoss() | |
| self.discriminator_losses = [] | |
| self.generator_losses = [] | |
| self.curr_step = 0 | |
| self.bce = nn.BCEWithLogitsLoss() | |
| self.l1_loss = nn.L1Loss() | |
| self.save_hyperparameters() | |
| def configure_optimizers(self): | |
| optimizer_G = optim.Adam(self.gen.parameters(), lr=self.hparams.learning_rate, betas=(0.5, 0.999)) | |
| optimizer_D = optim.Adam(self.disc.parameters(), lr=self.hparams.learning_rate, betas=(0.5, 0.999)) | |
| return optimizer_G, optimizer_D | |
| def on_load_checkpoint(self, checkpoint): | |
| # List of keys that you expect to load from the checkpoint | |
| keys_to_load = ['discriminator_losses', 'generator_losses', 'curr_step'] | |
| # Iterate over the keys and load them if they exist in the checkpoint | |
| for key in keys_to_load: | |
| if key in checkpoint: | |
| setattr(self, key, checkpoint[key]) | |
| def on_save_checkpoint(self, checkpoint): | |
| # Save the current state of the model | |
| checkpoint['discriminator_losses'] = self.discriminator_losses | |
| checkpoint['generator_losses'] = self.generator_losses | |
| checkpoint['curr_step'] = self.curr_step | |
| def training_step(self, batch, batch_idx): | |
| # Get the Optimizers | |
| opt_generator, opt_discriminator = self.optimizers() | |
| X, y = batch | |
| # Train Discriminator | |
| y_fake = self.gen(X) | |
| D_real = self.disc(X, y) | |
| D_fake = self.disc(X, y_fake.detach()) | |
| D_real_loss = self.loss_fn(D_real, torch.ones_like(D_real)) | |
| D_fake_loss = self.loss_fn(D_fake, torch.zeros_like(D_fake)) | |
| D_loss = (D_real_loss + D_fake_loss) / 2 | |
| opt_discriminator.zero_grad() | |
| self.manual_backward(D_loss) | |
| opt_discriminator.step() | |
| self.log("D_loss", D_loss.item(), on_step=False, on_epoch=True, prog_bar=True) | |
| self.discriminator_losses.append(D_loss.item()) | |
| # Train Generator | |
| D_fake = self.disc(X, y_fake) | |
| G_fake_loss = self.bce(D_fake, torch.ones_like(D_fake)) | |
| L1 = self.l1_loss(y_fake, y) * self.hparams.l1_lambda | |
| G_loss = G_fake_loss + L1 | |
| opt_generator.zero_grad() | |
| self.manual_backward(G_loss) | |
| opt_generator.step() | |
| self.log("G_loss", G_loss.item(), on_step=False, on_epoch=True, prog_bar=True) | |
| self.generator_losses.append(G_loss.item()) | |
| self.log("Current_Step", self.curr_step, on_step=False, on_epoch=True, prog_bar=True) | |
| # Visualize | |
| if self.curr_step % self.hparams.display_step == 0 and self.curr_step > 0: | |
| save_some_examples(self.gen, batch, self.current_epoch) | |
| self.curr_step += 1 | |