Update dataloder_pytorch.py

dataloder_pytorch.py

@@ -46,3 +46,63 @@ val_data = DataLoader(CarShadowDataset(root_dir='dataset/val', transform=your_tr
 for car_image, shadow_image in train_data:
   # Access your data for training
   # ...
+
+
+# ... (Previous code for model definition and DataLoader)
+
+# Discriminator Training
+def train_discriminator(d_optimizer, real_images, fake_images, real_labels, fake_labels):
+  # Clear gradients
+  d_optimizer.zero_grad()
+
+  # Forward pass through discriminator
+  d_real_output = discriminator(real_images, real_images)  # Real images with real shadows
+  d_fake_output = discriminator(real_images, fake_images)  # Real images with generated shadows
+
+  # Calculate loss
+  d_real_loss = criterion(d_real_output, torch.ones_like(d_real_output))
+  d_fake_loss = criterion(d_fake_output, torch.zeros_like(d_fake_output))
+  d_loss = (d_real_loss + d_fake_loss) / 2
+
+  # Backpropagate and update weights
+  d_loss.backward()
+  d_optimizer.step()
+
+  # Return the discriminator loss
+  return d_loss.item()
+
+# Generator Training
+def train_generator(g_optimizer, real_images, fake_images):
+  # Clear gradients
+  g_optimizer.zero_grad()
+
+  # Forward pass through discriminator (using generated shadows)
+  g_fake_output = discriminator(real_images, fake_images)
+
+  # Calculate loss (try to fool the discriminator)
+  g_loss = criterion(g_fake_output, torch.ones_like(g_fake_output))
+
+  # Backpropagate and update weights
+  g_loss.backward()
+  g_optimizer.step()
+
+  # Return the generator loss
+  return g_loss.item()
+
+# Training loop
+for epoch in range(epochs):
+  for i, (real_images, real_shadows) in enumerate(train_data):
+    # Generate fake shadows
+    fake_shadows = generator(real_images)
+
+    # Train discriminator
+    d_loss = train_discriminator(d_optimizer, real_images, fake_shadows, torch.ones(real_images.size(0)), torch.zeros(real_images.size(0)))
+
+    # Train generator
+    g_loss = train_generator(g_optimizer, real_images, fake_shadows)
+
+    # Print training progress
+    if i % 100 == 0:
+      print(f'Epoch [{epoch+1}/{epochs}], Step [{i+1}/{len(train_data)}], D_loss: {d_loss:.4f}, G_loss: {g_loss:.4f}')
+
+