misc: resolve some review comments for product consistency

Author: samsara-ku

examples/pytorch/domain_templates/generative_adversarial_net_ddp.py

@@ -20,19 +20,12 @@
 """
 
 import math
-
-# ! TESTING
-import os
-import sys
 from argparse import ArgumentParser, Namespace
 
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
 
-sys.path.append(os.path.join(os.getcwd(), "src"))
-# ! TESTING
-
 from lightning.pytorch import cli_lightning_logo
 from lightning.pytorch.core import LightningModule
 from lightning.pytorch.demos.mnist_datamodule import MNISTDataModule
@@ -44,7 +37,7 @@
     import torchvision
 
 
-def _block(in_feat: int, out_feat: int, normalize: bool = True):
+def _block(in_feat: int, out_feat: int, normalize: bool = True) -> list:
     layers = [nn.Linear(in_feat, out_feat)]
     if normalize:
         layers.append(nn.BatchNorm1d(out_feat, 0.8))
@@ -135,10 +128,6 @@ def __init__(
 
         self.example_input_array = torch.zeros(2, self.hparams.latent_dim)
 
-        # ! TESTING
-        self.save_path = "pl_test_multi_gpu"
-        os.makedirs(self.save_path, exist_ok=True)
-
     def forward(self, z):
         return self.generator(z)
 
@@ -203,36 +192,25 @@ def configure_optimizers(self):
         opt_d = torch.optim.Adam(self.discriminator.parameters(), lr=lr, betas=(b1, b2))
         return opt_g, opt_d
 
-    # ! TESTING
-    def on_train_epoch_start(self):
-        if self.trainer.is_global_zero:
-            print("GEN: ", self.generator.module.model[0].bias[:10])
-            print("DISC: ", self.discriminator.module.model[0].bias[:10])
-
-    # ! TESTING
-    def validation_step(self, batch, batch_idx):
-        pass
+    def on_train_epoch_end(self):
+        z = self.validation_z.type_as(self.generator.model[0].weight)
 
-    # ! TESTING
-    @torch.no_grad()
-    def on_validation_epoch_end(self):
-        if not self.current_epoch % 5:
-            return
-        self.generator.eval(), self.discriminator.eval()
-
-        z = self.validation_z.type_as(self.generator.module.model[0].weight)
+        # log sampled images`
         sample_imgs = self(z)
-
-        if self.trainer.is_global_zero:
-            grid = torchvision.utils.make_grid(sample_imgs)
-            torchvision.utils.save_image(grid, os.path.join(self.save_path, f"epoch_{self.current_epoch}.png"))
-
-        self.generator.train(), self.discriminator.train()
+        grid = torchvision.utils.make_grid(sample_imgs)
+        for logger in self.loggers:
+            logger.experiment.add_image("generated_images", grid, self.current_epoch)
 
 
 def main(args: Namespace) -> None:
+    # ------------------------
+    # 1 INIT LIGHTNING MODEL
+    # ------------------------
     model = GAN(lr=args.lr, b1=args.b1, b2=args.b2, latent_dim=args.latent_dim)
 
+    # ------------------------
+    # 2 INIT TRAINER
+    # ------------------------
     # ! `MultiModelDDPStrategy` is critical for multi-gpu training
     # ! Otherwise, it will not work with multiple models.
     # ! There are two ways to run training codes with previous `DDPStrategy`;
@@ -246,6 +224,9 @@ def main(args: Namespace) -> None:
         max_epochs=100,
     )
 
+    # ------------------------
+    # 3 START TRAINING
+    # ------------------------
     trainer.fit(model, dm)
 
 

src/lightning/pytorch/strategies/ddp.py

@@ -421,7 +421,7 @@ def teardown(self) -> None:
 
 class MultiModelDDPStrategy(DDPStrategy):
     @override
-    def _setup_model(self, model: Module) -> Module:
+    def _setup_model(self, model: Module) -> DistributedDataParallel:
         device_ids = self.determine_ddp_device_ids()
         log.debug(f"setting up DDP model with device ids: {device_ids}, kwargs: {self._ddp_kwargs}")
         # https://pytorch.org/docs/stable/notes/cuda.html#id5