pytorch lightning fixes

This reverts commit 02d9bb6.

Author: bcebere

decaf/DECAF.py

@@ -346,76 +346,66 @@ def get_gen_order(self) -> list:
         return gen_order
 
     def training_step(
-        self, batch: torch.Tensor, batch_idx: int, **kwargs
+        self, batch: torch.Tensor, batch_idx: int, optimizer_idx: int
     ) -> OrderedDict:
-        optimizer_g, optimizer_d = self.optimizers()
-
         # sample noise
         z = self.sample_z(batch.shape[0])
         z = z.type_as(batch)
         generated_batch = self.generator.sequential(batch, z, self.get_gen_order())
 
-        # train disc
-        #
-        self.toggle_optimizer(optimizer_d)
-        self.iterations_d += 1
-        # Measure discriminator's ability to classify real from generated samples
+        # train generator
+        if optimizer_idx == 0:
+            self.iterations_d += 1
+            # Measure discriminator's ability to classify real from generated samples
 
-        # how well can it label as real?
-        real_loss = torch.mean(self.discriminator(batch))
-        fake_loss = torch.mean(self.discriminator(generated_batch.detach()))
+            # how well can it label as real?
+            real_loss = torch.mean(self.discriminator(batch))
+            fake_loss = torch.mean(self.discriminator(generated_batch.detach()))
 
-        # discriminator loss
-        d_loss = fake_loss - real_loss
+            # discriminator loss
+            d_loss = fake_loss - real_loss
 
-        # add the gradient penalty
-        d_loss += self.hparams.lambda_gp * self.compute_gradient_penalty(
-            batch, generated_batch
-        )
-        if torch.isnan(d_loss).sum() != 0:
-            raise ValueError("NaN in the discr loss")
-
-        self.manual_backward(d_loss)
-        optimizer_d.step()
-        optimizer_d.zero_grad()
-        self.untoggle_optimizer(optimizer_d)
-        # train gen
-        #
-        # sanity check: keep track of G updates
-        self.toggle_optimizer(optimizer_g)
-        self.iterations_g += 1
-
-        # adversarial loss (negative D fake loss)
-        generated_batch = self.generator.sequential(batch, z, self.get_gen_order())
-        g_loss = -torch.mean(
-            self.discriminator(generated_batch)
-        )  # self.adversarial_loss(self.discriminator(self.generated_batch), valid)
+            # add the gradient penalty
+            d_loss += self.hparams.lambda_gp * self.compute_gradient_penalty(
+                batch, generated_batch
+            )
+            if torch.isnan(d_loss).sum() != 0:
+                raise ValueError("NaN in the discr loss")
+
+            return d_loss
+        elif optimizer_idx == 1:
+            # sanity check: keep track of G updates
+            self.iterations_g += 1
+
+            # adversarial loss (negative D fake loss)
+            g_loss = -torch.mean(
+                self.discriminator(generated_batch)
+            )  # self.adversarial_loss(self.discriminator(self.generated_batch), valid)
+
+            # add privacy loss of ADS-GAN
+            g_loss += self.hparams.lambda_privacy * self.privacy_loss(
+                batch, generated_batch
+            )
 
-        # add privacy loss of ADS-GAN
-        g_loss += self.hparams.lambda_privacy * self.privacy_loss(
-            batch, generated_batch
-        )
+            # add l1 regularization loss
+            g_loss += self.hparams.l1_g * self.l1_reg(self.generator)
 
-        # add l1 regularization loss
-        g_loss += self.hparams.l1_g * self.l1_reg(self.generator)
+            if len(self.dag_seed) == 0:
+                if self.hparams.grad_dag_loss:
+                    g_loss += self.gradient_dag_loss(batch, z)
+            if torch.isnan(g_loss).sum() != 0:
+                raise ValueError("NaN in the gen loss")
 
-        if len(self.dag_seed) == 0:
-            if self.hparams.grad_dag_loss:
-                g_loss += self.gradient_dag_loss(batch, z)
-        if torch.isnan(g_loss).sum() != 0:
-            raise ValueError("NaN in the gen loss")
-        self.manual_backward(g_loss)
-        optimizer_g.step()
-        optimizer_g.zero_grad()
-        self.untoggle_optimizer(optimizer_g)
+            return g_loss
+        else:
+            raise ValueError("should not get here")
 
     def configure_optimizers(self) -> tuple:
         lr = self.hparams.lr
         b1 = self.hparams.b1
         b2 = self.hparams.b2
         weight_decay = self.hparams.weight_decay
 
-        self.automatic_optimization = False
         opt_g = torch.optim.AdamW(
             self.generator.parameters(),
             lr=lr,

requirements.txt

@@ -2,7 +2,7 @@ loguru
 networkx>=2.0,<3.0
 numpy>=1.19
 pandas
-pytorch-lightning>=1.4
+pytorch-lightning<2.0
 scikit-learn
 scipy
 torch>=1.9