Upload discriminator.py

discriminator.py

@@ -0,0 +1,37 @@
+import numpy as np
+
+import torch.nn as nn
+import torch
+
+
+class Discriminator(nn.Module):
+    def __init__(
+        self,
+        image_shape: (int, int, int),
+        use_cuda: bool = False,
+        saved_model: str or None = None
+    ):
+        super(Discriminator, self).__init__()
+
+        self.model = nn.Sequential(
+            nn.Linear(int(np.prod(image_shape)), 512),
+            nn.LeakyReLU(0.2, inplace=True),
+            nn.Linear(512, 256),
+            nn.LeakyReLU(0.2, inplace=True),
+            nn.Linear(256, 1),
+        )
+        if saved_model is not None:
+            self.model.load_state_dict(
+                torch.load(
+                    saved_model,
+                    map_location=torch.device('cuda' if use_cuda else 'cpu')
+                )
+            )
+
+    def forward(self, img):
+        img_flat = img.view(img.shape[0], -1)
+        validity = self.model(img_flat)
+        return validity
+
+    def save(self, to):
+        torch.save(self.model.state_dict(), to)