updates

src/geo_veritas/__init__.py

@@ -0,0 +1,4 @@
+from .datamodule import MNISTDataModule
+from .generator import Generator
+from .discriminator import Discriminator
+from .gan import GAN

src/geo_veritas/__main__.py

@@ -1,2 +1,13 @@
+from geo_veritas import MNISTDataModule
+from geo_veritas import GAN
+import lightning as L
+
 if __name__ == '__main__':
-    print("Hello, World!")
+    dm = MNISTDataModule()
+    model = GAN(*dm.dims)
+    trainer = L.Trainer(
+        accelerator="auto",
+        devices=1,
+        max_epochs=5,
+    )
+    trainer.fit(model, dm)

src/geo_veritas/gan.py

@@ -9,6 +9,9 @@
 import torchvision.transforms as transforms
 from torch.utils.data import DataLoader, random_split
 from torchvision.datasets import MNIST
+from geo_veritas import Generator, Discriminator
+
+BATCH_SIZE = 256 if torch.cuda.is_available() else 64
 
 class GAN(L.LightningModule):
     def __init__(
@@ -59,7 +62,7 @@ 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("generated_images", grid, 0)
 
         # ground truth result (ie: all fake)
         # put on GPU because we created this tensor inside training_loop
@@ -113,4 +116,4 @@ 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("generated_images", grid, self.current_epoch)