Update train.py

train.py

@@ -0,0 +1,84 @@
+import torch
+from torch.optim import Adam
+from torch.utils.data import Dataset, DataLoader
+from torchvision import transforms
+import os
+from PIL import Image
+import pickle
+
+# Load pre-trained model
+with open('ffhq.pkl', 'rb') as f:
+    data = pickle.load(f)
+
+G = data['G_ema']
+D = data['D']
+
+# Check if CUDA is available
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+G = G.to(device)
+D = D.to(device)
+
+# Custom dataset class
+class CustomDataset(Dataset):
+    def __init__(self, image_dir, transform=None):
+        self.image_dir = image_dir
+        self.image_files = [f for f in os.listdir(image_dir) if f.endswith('.jpg') or f.endswith('.png')]
+        self.transform = transform
+
+    def __len__(self):
+        return len(self.image_files)
+
+    def __getitem__(self, idx):
+        img_path = os.path.join(self.image_dir, self.image_files[idx])
+        image = Image.open(img_path).convert('RGB')
+        if self.transform:
+            image = self.transform(image)
+        return image
+
+# Data loading
+transform = transforms.Compose([
+    transforms.Resize((G.img_resolution, G.img_resolution)),
+    transforms.ToTensor(),
+    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
+])
+
+dataset = CustomDataset("/path/to/your/image_dir", transform=transform)
+dataloader = DataLoader(dataset, batch_size=4, shuffle=True)
+
+# Fine-tuning setup
+optimizer_g = Adam(G.parameters(), lr=0.0001, betas=(0, 0.99))
+optimizer_d = Adam(D.parameters(), lr=0.0001, betas=(0, 0.99))
+scheduler_g = torch.optim.lr_scheduler.ExponentialLR(optimizer_g, gamma=0.99)
+scheduler_d = torch.optim.lr_scheduler.ExponentialLR(optimizer_d, gamma=0.99)
+
+num_epochs = 100
+for epoch in range(num_epochs):
+    for batch in dataloader:
+        real_images = batch.to(device)
+
+        # Generate fake images
+        z = torch.randn([batch.size(0), G.z_dim]).to(device)
+        fake_images = G(z, None)
+
+        # Compute losses
+        g_loss = -torch.mean(torch.log(D(fake_images, None)))
+        d_loss_real = -torch.mean(torch.log(D(real_images, None)))
+        d_loss_fake = -torch.mean(torch.log(1 - D(fake_images, None)))
+        d_loss = d_loss_real + d_loss_fake
+
+        # Update models
+        optimizer_g.zero_grad()
+        g_loss.backward()
+        optimizer_g.step()
+
+        optimizer_d.zero_grad()
+        d_loss.backward()
+        optimizer_d.step()
+
+    scheduler_g.step()
+    scheduler_d.step()
+
+    print(f"Epoch {epoch+1}/{num_epochs}, G Loss: {g_loss.item()}, D Loss: {d_loss.item()}")
+
+# Save the fine-tuned model
+torch.save(G.state_dict(), 'fine_tuned_stylegan.pth')