sample and train.py implementations

sample.py

@@ -0,0 +1,9 @@
+from vegans import Vegans
+
+vegans = Vegans
+
+vegans.load_pretrained()
+
+text = 'something'
+
+vegans.generate(text)

train.py

@@ -0,0 +1,11 @@
+from vegans import Vegans
+
+# To Start training loop, it's simple
+
+vegans = Vegans()
+
+vegans.train()
+
+# or you can simply just
+# vegans.load_pretrained()
+

vegans/__init__.py

@@ -0,0 +1 @@
+from .vegans import Vegans

vegans/dataset.py

@@ -5,14 +5,22 @@ from datasets import load_dataset
 
 
 class TrainingSet(Dataset):
-    def __init__(self,dataset_name='',image_size=(128,128)):
+    def __init__(self,image_size=(128,128),gray_scale=True):
         super().__init__()        
         self.transforms = tran.Compose([
-
+            tran.Resize(image_size),
+            tran.ToTensor(),
+            tran.Normalize([0.5],[0.5]),
+            tran.Grayscale() if gray_scale==True else None,
         ])
+
+        self.images = None
+        self.labels = None
+ 
     def __getitem__(self, index):
-        return self
+        image = self.transforms(self.images[index])
+        label = self.labels[index]
+        return image,label
     
     def __len__(self):
-        return self
-    
+        return len(self.images)

vegans/discriminator.py

@@ -5,6 +5,9 @@ import torch
 class Discriminator(nn.Module):
   def __init__(self):
     super(Discriminator, self).__init__()
+    self.model = nn.Sequential(
 
-  def forward(self):
-    pass
+    )
+
+  def forward(self,x,y):
+    return x

vegans/generator.py

@@ -5,7 +5,7 @@ import torch
 class Generator(nn.Module):
   def __init__(self):
     super(Generator, self).__init__()
-    def block():
+    def block(in_feat,out_feat,norm=False):
       pass
     
     self.model = nn.Sequential(
@@ -14,6 +14,3 @@ class Generator(nn.Module):
 
   def forward(self):
     pass
-
-  def generate(self,text):
-    pass

vegans/utils.py

@@ -1 +1,10 @@
+def display_image():
+    pass
 
+def save_image(image,location):
+    pass
+
+def log(text):
+    print("#############################################\n")
+    print(f"{text}\n")
+    print("#############################################\n")

vegans/vegans.py

@@ -3,9 +3,10 @@ import torch
 from generator import Generator
 from discriminator import Discriminator
 from discriminator import Discriminator
-from utils import save_image,display_image
+from utils import save_image,display_image,log
 from dataset import TrainingSet
 
+
 class Vegans:
     def __init__(self,lr=0.01,b1=0.02,b2=0.02):
         self.learning_rate = lr
@@ -19,15 +20,35 @@ class Vegans:
         self.d_optim = torch.optim.Adam(self.discriminator.parameters(),lr=self.learning_rate,betas=self.betas)
 
     def train(self):
+        loss_fn = torch.nn.BCELoss()
+        log("Started Training Loop")
+
         for epoch in range(self.num_epoch):
             for i,(image,label) in enumerate(self.dataset):
+                
                 self.g_optim.zero_grad()
+
+                loss = loss_fn(0,0)
+
+
                 self.d_optim.zero_grad()
 
+            print(f"Epoch {epoch} done. Loss is {loss.item()}`")
 
+        log("Finish Training")
 
     def generate(self,label):
-        pass
+        # TODO
+        noise = 0
+        output = self.generator(noise,label)
 
     def load_pretrained(self,checkpoint):
-        pass
+
+        # TODO
+        saved_checkpoint = checkpoint
+
+        model_checkpoint = torch.load('')
+        self.generator.load_state_dict(model_checkpoint)
+
+        log("Successfully loaded model")
+