Create train.py

train.py

@@ -0,0 +1,33 @@
+from model import ConditionalGenerator, Discriminator
+from data_loader import MultiStyleDataset
+import torch.optim as optim
+
+
+num_styles = len(os.listdir(".styles/"))  
+G = ConditionalGenerator(num_styles)
+D = Discriminator(num_styles)
+
+opt_G = optim.Adam(G.parameters(), lr=2e-4)
+opt_D = optim.Adam(D.parameters(), lr=2e-4)
+
+dataset = MultiStyleDataset(".styles/")
+
+for epoch in range(100):
+    for img, style_id in dataset:
+        
+        fake_img = G(img.unsqueeze(0), torch.tensor([style_id]))
+        
+
+        real_loss = torch.mean((D(img.unsqueeze(0), torch.tensor([style_id])) - 1)**2)
+        fake_loss = torch.mean(D(fake_img.detach(), torch.tensor([style_id]))**2)
+        loss_D = (real_loss + fake_loss) / 2
+        
+        opt_D.zero_grad()
+        loss_D.backward()
+        opt_D.step()
+        
+        
+        loss_G = torch.mean((D(fake_img, torch.tensor([style_id])) - 1)**2
+        opt_G.zero_grad()
+        loss_G.backward()
+        opt_G.step()