Create train.py

train.py

@@ -0,0 +1,243 @@
+import torch
+import torch.nn as nn
+import torch.optim as optim
+import torchvision.transforms as transforms
+from torch.utils.data import DataLoader, Dataset
+from datasets import load_dataset
+import matplotlib
+matplotlib.use('Agg')
+import matplotlib.pyplot as plt
+import numpy as np
+from torch.cuda.amp import autocast, GradScaler
+import torchvision.utils as vutils
+from IPython.display import display
+
+# --- FaceGen v1 Config ---
+BATCH_SIZE = 128
+IMAGE_SIZE = 128
+CHANNELS = 3
+Z_DIM = 128
+FEATURES_G = 256
+FEATURES_D = 128
+EPOCHS = 250
+LR = 0.0002
+BETA1 = 0.5
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+print(f"Training will run on: {device}")
+
+print("Loading face dataset...")
+hf_dataset = load_dataset("SDbiaseval/faces", split="train")
+
+transform = transforms.Compose([
+    transforms.Resize(IMAGE_SIZE),
+    transforms.CenterCrop(IMAGE_SIZE),
+    transforms.RandomHorizontalFlip(),
+    transforms.ToTensor(),
+    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
+])
+
+class FaceDataset(Dataset):
+    def __init__(self, hf_ds, transform):
+        self.hf_ds = hf_ds
+        self.transform = transform
+    def __len__(self):
+        return len(self.hf_ds)
+    def __getitem__(self, idx):
+        img = self.hf_ds[idx]['image'].convert("RGB")
+        return self.transform(img)
+
+dataset = FaceDataset(hf_dataset, transform)
+
+dataloader = DataLoader(
+    dataset,
+    batch_size=BATCH_SIZE,
+    shuffle=True,
+    drop_last=True,
+    num_workers=4,
+    pin_memory=True
+)
+print(f"Dataset ready with {len(dataset)} faces.")
+
+class Generator(nn.Module):
+    def __init__(self, z_dim, channels, features_g):
+        super(Generator, self).__init__()
+        self.net = nn.Sequential(
+            # Input: Z_DIM x 1 x 1
+            nn.ConvTranspose2d(z_dim, features_g * 16, 4, 1, 0, bias=False),
+            nn.BatchNorm2d(features_g * 16),
+            nn.ReLU(True),
+            # 4x4 -> 8x8
+            nn.ConvTranspose2d(features_g * 16, features_g * 8, 4, 2, 1, bias=False),
+            nn.BatchNorm2d(features_g * 8),
+            nn.ReLU(True),
+            # 8x8 -> 16x16
+            nn.ConvTranspose2d(features_g * 8, features_g * 4, 4, 2, 1, bias=False),
+            nn.BatchNorm2d(features_g * 4),
+            nn.ReLU(True),
+            # 16x16 -> 32x32
+            nn.ConvTranspose2d(features_g * 4, features_g * 2, 4, 2, 1, bias=False),
+            nn.BatchNorm2d(features_g * 2),
+            nn.ReLU(True),
+            # 32x32 -> 64x64
+            nn.ConvTranspose2d(features_g * 2, features_g, 4, 2, 1, bias=False),
+            nn.BatchNorm2d(features_g),
+            nn.ReLU(True),
+            # 64x64 -> 128x128
+            nn.ConvTranspose2d(features_g, channels, 4, 2, 1, bias=False),
+            nn.Tanh()
+        )
+
+    def forward(self, x):
+        return self.net(x)
+
+netG = Generator(Z_DIM, CHANNELS, FEATURES_G).to(device)
+
+class Discriminator(nn.Module):
+    def __init__(self, channels, features_d):
+        super(Discriminator, self).__init__()
+        self.net = nn.Sequential(
+            # 128x128 -> 64x64
+            nn.Conv2d(channels, features_d, 4, 2, 1, bias=False),
+            nn.LeakyReLU(0.2, inplace=True),
+            # 64x64 -> 32x32
+            nn.Conv2d(features_d, features_d * 2, 4, 2, 1, bias=False),
+            nn.BatchNorm2d(features_d * 2),
+            nn.LeakyReLU(0.2, inplace=True),
+            # 32x32 -> 16x16
+            nn.Conv2d(features_d * 2, features_d * 4, 4, 2, 1, bias=False),
+            nn.BatchNorm2d(features_d * 4),
+            nn.LeakyReLU(0.2, inplace=True),
+            # 16x16 -> 8x8
+            nn.Conv2d(features_d * 4, features_d * 8, 4, 2, 1, bias=False),
+            nn.BatchNorm2d(features_d * 8),
+            nn.LeakyReLU(0.2, inplace=True),
+            # 8x8 -> 4x4
+            nn.Conv2d(features_d * 8, features_d * 16, 4, 2, 1, bias=False),
+            nn.BatchNorm2d(features_d * 16),
+            nn.LeakyReLU(0.2, inplace=True),
+            # 4x4 -> 1x1
+            nn.Conv2d(features_d * 16, 1, 4, 1, 0, bias=False),
+        )
+
+    def forward(self, x):
+        return self.net(x)
+
+netD = Discriminator(CHANNELS, FEATURES_D).to(device)
+
+def weights_init(m):
+    classname = m.__class__.__name__
+    if classname.find('Conv') != -1:
+        nn.init.normal_(m.weight.data, 0.0, 0.02)
+    elif classname.find('BatchNorm') != -1:
+        nn.init.normal_(m.weight.data, 1.0, 0.02)
+        nn.init.constant_(m.bias.data, 0)
+
+netG.apply(weights_init)
+netD.apply(weights_init)
+
+criterion = nn.BCEWithLogitsLoss()
+
+optG = optim.Adam(netG.parameters(), lr=LR, betas=(BETA1, 0.999))
+optD = optim.Adam(netD.parameters(), lr=LR, betas=(BETA1, 0.999))
+
+fixed_noise = torch.randn(64, Z_DIM, 1, 1, device=device)
+
+scaler = torch.amp.GradScaler('cuda') 
+
+print(f"Model size G: {sum(p.numel() for p in netG.parameters())/1e6:.2f}M parameters")
+print(f"Model size D: {sum(p.numel() for p in netD.parameters())/1e6:.2f}M parameters")
+
+real_label_val = 0.9
+fake_label_val = 0.1
+
+for epoch in range(EPOCHS):
+    for i, real_images in enumerate(dataloader):
+        real_images = real_images.to(device)
+        b_size = real_images.size(0)
+        
+        # --- Discriminator Update ---
+        optD.zero_grad()
+        with torch.amp.autocast('cuda'):
+            output_real = netD(real_images).view(-1)
+            lossD_real = criterion(output_real, torch.full((b_size,), real_label_val, device=device))
+            
+            noise = torch.randn(b_size, Z_DIM, 1, 1, device=device)
+            fake_images = netG(noise)
+            output_fake = netD(fake_images.detach()).view(-1)
+            lossD_fake = criterion(output_fake, torch.full((b_size,), fake_label_val, device=device))
+            lossD = lossD_real + lossD_fake
+
+        scaler.scale(lossD).backward()
+        scaler.step(optD)
+        
+        # --- Generator Update ---
+        optG.zero_grad()
+        with torch.amp.autocast('cuda'):
+            output_fake_G = netD(fake_images).view(-1)
+            lossG = criterion(output_fake_G, torch.full((b_size,), real_label_val, device=device))
+
+        scaler.scale(lossG).backward()
+        scaler.step(optG)
+        scaler.update()
+
+        if i % 10 == 0:
+            print(f"E[{epoch}] I[{i}/{len(dataloader)}] Loss_D: {lossD.item():.4f} Loss_G: {lossG.item():.4f}")
+
+    if (epoch + 1) % 10 == 0 or epoch == 0:
+        netG.eval()
+        with torch.no_grad():
+            with torch.amp.autocast('cuda'):
+                sample = netG(fixed_noise[0:1]).detach().cpu().float()
+            
+            vutils.save_image(sample, f"face_sample_epoch_{epoch}.png", normalize=True)
+            print(f"--> Sample saved: face_sample_epoch_{epoch}.png")
+            
+        netG.train()
+
+    if (epoch + 1) % 50 == 0:
+        torch.save({
+            'epoch': epoch,
+            'model_state_dict': netG.state_dict(),
+            'optimizer_state_dict': optG.state_dict(),
+            'netD_state_dict': netD.state_dict(),
+            'optD_state_dict': optD.state_dict(),
+            'scaler_state_dict': scaler.state_dict(),
+        }, f'facegen_v1_checkpoint_epoch_{epoch+1}.ckpt')
+        print(f"--> Sicherheits-Checkpoint gespeichert: Epoche {epoch+1}")
+
+torch.save({
+    'epoch': EPOCHS,
+    'model_state_dict': netG.state_dict(),
+    'optimizer_state_dict': optG.state_dict(),
+    'netD_state_dict': netD.state_dict(),
+    'optD_state_dict': optD.state_dict(),
+    'scaler_state_dict': scaler.state_dict(),
+}, 'facegen_v1_full_checkpoint.ckpt')
+
+torch.save(netG.state_dict(), 'facegen_v1_generator_only.pth')
+
+print("Files saved: Training finished.")
+
+print("Doing professionell gallery export...")
+
+# --- FaceGen v2: Professional Gallery Export (Fix) ---
+netG.eval()
+
+with torch.no_grad():
+    with torch.amp.autocast('cuda'):
+        fake_faces = netG(fixed_noise).detach().cpu().float() 
+
+grid = vutils.make_grid(fake_faces, padding=4, normalize=True)
+grid_np = grid.numpy().transpose((1, 2, 0))
+
+plt.figure(figsize=(12, 12), facecolor='#111111') 
+plt.imshow(grid_np, interpolation='bilinear')
+plt.axis("off")
+
+plt.title(f"FaceGen v1 | Training Complete | {FEATURES_G}x{FEATURES_D} Filters", 
+          color='white', fontsize=16, fontweight='bold', pad=20)
+
+plt.tight_layout()
+
+plt.savefig("facegen_v2_results.png", facecolor='#111111', bbox_inches='tight')