Cleanup

.gitignore

@@ -126,4 +126,4 @@ venv.bak/
 dmypy.json
 
 # Pyre type checker
-.pyre/
+.pyre/

code_soup/ch5/__init__.py

@@ -1 +1,2 @@
-from code_soup.ch5.dummy import add_nums
+from code_soup.ch5.datasets import MnistDataset
+from code_soup.ch5.models import Discriminator, Generator

code_soup/ch5/datasets/__init__.py

@@ -0,0 +1 @@
+from code_soup.ch5.datasets.mnist import MnistDataset

code_soup/chapter_5/datasets/mnist.py → code_soup/ch5/datasets/mnist.py

@@ -1,8 +1,8 @@
-from torch.utils.data import Dataset
 import torchvision.datasets as datasets
+from torch.utils.data import Dataset
 
 
-class Mnist(Dataset):
+class MnistDataset(Dataset):
     def __init__(self, transform=None):
 
         self.train_data = datasets.MNIST(

code_soup/chapter_5/mnist_gan.py → code_soup/ch5/mnist_gan.py

@@ -1,10 +1,12 @@
-from datasets import Mnist
-from models import Generator, Discriminator
+import argparse
+
+import torch
+import torch.optim as optim
 import torchvision.transforms as transforms
 import torchvision.utils as vutils
-import torch.optim as optim
-import torch
-import argparse
+
+from code_soup.ch5.datasets import MnistDataset
+from code_soup.ch5.models import Discriminator, Generator
 
 parser = argparse.ArgumentParser(
     prog="mnist_gan.py", description="Train an MNIST GAN model"
@@ -52,7 +54,7 @@ def train_mnist_gan():
             transforms.Normalize((0.5,), (0.5,)),
         ]
     )
-    dataset = Mnist(transform=transform)
+    dataset = MnistDataset(transform=transform)
     dataloader = torch.utils.data.DataLoader(
         dataset, batch_size=dataloader_batch_size, shuffle=True
     )
@@ -92,7 +94,7 @@ def train_mnist_gan():
             errD_real.backward()
 
             D_x = output.mean().item()
-            ## Train with all-fake batch
+            # Train with all-fake batch
             # Generate batch of latent vectors
             noise = torch.randn(batch_size, latent_dims, device=device)
             # Generate fake image batch with G
@@ -137,9 +139,10 @@ def train_mnist_gan():
                         D_G_z2,
                     )
                 )
-        #save model weights
+        # save model weights
         torch.save(discriminator.state_dict(), "./discriminator.pth")
         torch.save(generator.state_dict(), "./generator.pth")
 
+
 if __name__ == "__main__":
     train_mnist_gan()

code_soup/ch5/models/__init__.py

@@ -0,0 +1,2 @@
+from code_soup.ch5.models.discriminator import Discriminator
+from code_soup.ch5.models.generator import Generator

requirements.txt

@@ -1,4 +1,4 @@
 numpy==1.21.1
 Pillow==8.3.1
 torch==1.9.0
-torchvision==0.10.0
+torchvision==0.10.0

tests/test_ch5/__init__.py

@@ -1 +1,2 @@
-from tests.test_ch5.test_dummy import test_add_nums
+from tests.test_ch5.test_datasets import TestMnistDataset
+from tests.test_ch5.test_models import TestDiscriminatorModel, TestGeneratorModel

tests/test_ch5/test_datasets/__init__.py

@@ -1 +1 @@
-from .mnist_test import TestMnistDataset
+from tests.test_ch5.test_datasets.mnist_test import TestMnistDataset

tests/test_ch5/test_datasets/mnist_test.py

@@ -2,17 +2,20 @@
 from pathlib import Path
 
 import torch
+import torchvision.datasets as datasets
 from PIL import Image
 from torch.utils.data import DataLoader, Dataset
 from torchvision import transforms
-import torchvision.datasets as datasets
-from code_soup.chapter_5.datasets import *
+
+from code_soup.ch5.datasets import MnistDataset
 
 
 class TestMnistDataset(unittest.TestCase):
     @classmethod
     def setUpClass(cls) -> None:
-        mnist_dataset = Mnist(transform=transforms.Compose([transforms.ToTensor()]))
+        mnist_dataset = MnistDataset(
+            transform=transforms.Compose([transforms.ToTensor()])
+        )
         mnist_dataloader = DataLoader(mnist_dataset, batch_size=64, shuffle=False)
         cls.samples = next(iter(mnist_dataloader))
 

tests/test_ch5/test_models/__init__.py

@@ -1,2 +1,2 @@
-from .discriminator_test import TestDiscriminatorModel
-from .generator_test import TestGeneratorModel
+from tests.test_ch5.test_models.discriminator_test import TestDiscriminatorModel
+from tests.test_ch5.test_models.generator_test import TestGeneratorModel

tests/test_ch5/test_models/discriminator_test.py

@@ -1,7 +1,9 @@
+import unittest
+
 import torch
 import torch.nn as nn
-import unittest
-from code_soup.chapter_5.models import Discriminator
+
+from code_soup.ch5.models import Discriminator
 
 
 class TestDiscriminatorModel(unittest.TestCase):

tests/test_ch5/test_models/generator_test.py

@@ -1,7 +1,9 @@
 import unittest
-import torch.nn as nn
+
 import torch
-from code_soup.chapter_5.models import Generator
+import torch.nn as nn
+
+from code_soup.ch5.models import Generator
 
 
 class TestGeneratorModel(unittest.TestCase):