diff --git a/README.md b/README.md index ec8f720..f1d9d22 100644 --- a/README.md +++ b/README.md @@ -43,6 +43,19 @@ Forked from https://github.com/LynnHo/DCGAN-LSGAN-WGAN-WGAN-GP-Tensorflow ... ``` + +### PyTorch DCGAN training + +We now provide a PyTorch re-implementation of the CelebA DCGAN baseline that mirrors the behaviour of `train_celeba_dcgan.py`. + +``` +python train_celeba_dcgan_torch.py \ + --data-root ./data/img_align_celeba/img_align_celeba \ + --epochs 50 \ + --batch-size 64 +``` + +Generated sample grids are written to `./sample_images_while_training/celeba_dcgan_torch/` and checkpoints to `./checkpoints/celeba_dcgan_torch/`. Run `python train_celeba_dcgan_torch.py --help` for the full list of configurable options. ## Tensorboard ``` tensorboard --logdir=./logs/sia/ diff --git a/models_64x64_torch.py b/models_64x64_torch.py new file mode 100644 index 0000000..70a05c3 --- /dev/null +++ b/models_64x64_torch.py @@ -0,0 +1,115 @@ +"""PyTorch implementation of the 64x64 generator and discriminator networks. + +This module mirrors the TensorFlow version in :mod:`models_64x64` but uses +PyTorch's ``nn.Module`` classes so it can be integrated with a PyTorch +training pipeline. The architecture follows the standard DCGAN design that +is used throughout the original project. It intentionally keeps the API +minimal so the modules can be reused in standalone scripts or notebooks. +""" +from __future__ import annotations + +from dataclasses import dataclass + +import torch +from torch import nn + + +@dataclass +class DCGANConfig: + """Configuration describing the DCGAN model layout.""" + + latent_dim: int = 100 + base_channels: int = 64 + image_channels: int = 3 + + +class Generator(nn.Module): + """Generator network that maps latent vectors to 64x64 RGB images.""" + + def __init__(self, config: DCGANConfig = DCGANConfig()) -> None: + super().__init__() + self.config = config + + self.project = nn.Sequential( + nn.Linear(config.latent_dim, 4 * 4 * config.base_channels * 8, bias=False), + nn.BatchNorm1d(config.base_channels * 8 * 4 * 4), + nn.ReLU(True), + ) + + self.net = nn.Sequential( + nn.ConvTranspose2d(config.base_channels * 8, config.base_channels * 4, 4, 2, 1, bias=False), + nn.BatchNorm2d(config.base_channels * 4), + nn.ReLU(True), + nn.ConvTranspose2d(config.base_channels * 4, config.base_channels * 2, 4, 2, 1, bias=False), + nn.BatchNorm2d(config.base_channels * 2), + nn.ReLU(True), + nn.ConvTranspose2d(config.base_channels * 2, config.base_channels, 4, 2, 1, bias=False), + nn.BatchNorm2d(config.base_channels), + nn.ReLU(True), + nn.ConvTranspose2d(config.base_channels, config.image_channels, 4, 2, 1, bias=False), + nn.Tanh(), + ) + + def forward(self, z: torch.Tensor) -> torch.Tensor: + if z.dim() != 2 or z.size(1) != self.config.latent_dim: + raise ValueError( + f"Expected latent vectors of shape (batch, {self.config.latent_dim}), " + f"got {tuple(z.shape)}" + ) + x = self.project(z) + x = x.view(z.size(0), self.config.base_channels * 8, 4, 4) + return self.net(x) + + +class Discriminator(nn.Module): + """Discriminator network that scores 64x64 RGB images.""" + + def __init__(self, config: DCGANConfig = DCGANConfig()) -> None: + super().__init__() + self.config = config + + def block(in_channels: int, out_channels: int, normalize: bool = True) -> nn.Sequential: + layers = [nn.Conv2d(in_channels, out_channels, 4, 2, 1, bias=False)] + if normalize: + layers.append(nn.BatchNorm2d(out_channels)) + layers.append(nn.LeakyReLU(0.2, inplace=True)) + return nn.Sequential(*layers) + + self.net = nn.Sequential( + block(config.image_channels, config.base_channels, normalize=False), + block(config.base_channels, config.base_channels * 2), + block(config.base_channels * 2, config.base_channels * 4), + block(config.base_channels * 4, config.base_channels * 8), + ) + + self.head = nn.Conv2d(config.base_channels * 8, 1, 4, 1, 0, bias=False) + + def forward(self, img: torch.Tensor) -> torch.Tensor: + if img.dim() != 4 or img.size(1) != self.config.image_channels: + raise ValueError( + f"Expected images of shape (batch, {self.config.image_channels}, 64, 64), " + f"got {tuple(img.shape)}" + ) + features = self.net(img) + logits = self.head(features).view(img.size(0)) + return logits + + +def weights_init(module: nn.Module) -> None: + """Initialize model weights following the original DCGAN recipe.""" + + if isinstance(module, (nn.Conv2d, nn.ConvTranspose2d, nn.Linear)): + nn.init.normal_(module.weight.data, 0.0, 0.02) + if getattr(module, "bias", None) is not None: + nn.init.constant_(module.bias.data, 0.0) + elif isinstance(module, (nn.BatchNorm1d, nn.BatchNorm2d)): + nn.init.normal_(module.weight.data, 1.0, 0.02) + nn.init.constant_(module.bias.data, 0.0) + + +__all__ = [ + "DCGANConfig", + "Generator", + "Discriminator", + "weights_init", +] diff --git a/train_celeba_dcgan_torch.py b/train_celeba_dcgan_torch.py new file mode 100644 index 0000000..f42c520 --- /dev/null +++ b/train_celeba_dcgan_torch.py @@ -0,0 +1,203 @@ +"""PyTorch training script for the CelebA DCGAN baseline. + +This script mirrors the behaviour of :mod:`train_celeba_dcgan` but uses the +PyTorch models defined in :mod:`models_64x64_torch`. The defaults are chosen +so that existing datasets can be reused without additional preprocessing. +""" +from __future__ import annotations + +import argparse +import math +from datetime import datetime +from pathlib import Path +from typing import Tuple + +import torch +from PIL import Image +from torch import nn, optim +from torch.utils.data import DataLoader, Dataset +from torchvision import transforms, utils as vutils + +from models_64x64_torch import DCGANConfig, Discriminator, Generator, weights_init + + +class ImageFolderFlat(Dataset): + """Load images from a directory without requiring class sub-folders.""" + + def __init__(self, root: str, extensions: Tuple[str, ...] = (".jpg", ".png"), transform=None) -> None: + self.root = Path(root) + self.transform = transform + self.paths = sorted( + path + for ext in extensions + for path in self.root.glob(f"**/*{ext}") + if path.is_file() + ) + if not self.paths: + raise RuntimeError( + f"No images with extensions {extensions} were found under {self.root}." + ) + + def __len__(self) -> int: + return len(self.paths) + + def __getitem__(self, idx: int) -> torch.Tensor: + path = self.paths[idx] + with Image.open(path) as img: + img = img.convert("RGB") + if self.transform is not None: + img = self.transform(img) + return img + + +def build_dataloader(data_root: str, batch_size: int, workers: int) -> DataLoader: + transform = transforms.Compose( + [ + transforms.CenterCrop(108), + transforms.Resize(64), + transforms.ToTensor(), + transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)), + ] + ) + dataset = ImageFolderFlat(data_root, transform=transform) + return DataLoader(dataset, batch_size=batch_size, shuffle=True, num_workers=workers, pin_memory=True) + + +def train(args: argparse.Namespace) -> None: + device = torch.device("cuda" if torch.cuda.is_available() and not args.no_cuda else "cpu") + if torch.backends.cudnn.is_available(): + torch.backends.cudnn.benchmark = True + + loader = build_dataloader(args.data_root, args.batch_size, args.workers) + config = DCGANConfig(latent_dim=args.latent_dim, base_channels=args.channels) + + generator = Generator(config).to(device) + discriminator = Discriminator(config).to(device) + + generator.apply(weights_init) + discriminator.apply(weights_init) + + criterion = nn.BCEWithLogitsLoss() + optim_g = optim.Adam(generator.parameters(), lr=args.lr, betas=(0.5, 0.999)) + optim_d = optim.Adam(discriminator.parameters(), lr=args.lr, betas=(0.5, 0.999)) + + fixed_noise = torch.randn(args.sample_grid ** 2, args.latent_dim, device=device) + + global_step = 0 + for epoch in range(args.epochs): + for batch_idx, real in enumerate(loader): + real = real.to(device) + bsz = real.size(0) + + # --- Train discriminator --- + optim_d.zero_grad(set_to_none=True) + noise = torch.randn(bsz, args.latent_dim, device=device) + fake = generator(noise).detach() + + logits_real = discriminator(real) + logits_fake = discriminator(fake) + + labels_real = torch.ones_like(logits_real) + labels_fake = torch.zeros_like(logits_fake) + + loss_real = criterion(logits_real, labels_real) + loss_fake = criterion(logits_fake, labels_fake) + loss_d = (loss_real + loss_fake) * 0.5 + loss_d.backward() + optim_d.step() + + # --- Train generator --- + optim_g.zero_grad(set_to_none=True) + noise = torch.randn(bsz, args.latent_dim, device=device) + fake = generator(noise) + logits_fake = discriminator(fake) + labels_for_g = torch.ones_like(logits_fake) + loss_g = criterion(logits_fake, labels_for_g) + loss_g.backward() + optim_g.step() + + if global_step % args.log_interval == 0: + print( + f"Epoch {epoch + 1:03d}/{args.epochs} | Batch {batch_idx + 1:04d}/{len(loader):04d} " + f"| D loss: {loss_d.item():.4f} | G loss: {loss_g.item():.4f}" + ) + + if global_step % args.sample_interval == 0: + save_samples(generator, fixed_noise, args.sample_dir, args.sample_grid, global_step) + + if global_step % args.checkpoint_interval == 0: + save_checkpoint( + generator, + discriminator, + optim_g, + optim_d, + args.checkpoint_dir, + epoch, + global_step, + ) + + global_step += 1 + + print("Training finished.") + + +def save_samples( + generator: Generator, noise: torch.Tensor, directory: str, grid_size: int, step: int +) -> None: + directory_path = Path(directory) + directory_path.mkdir(parents=True, exist_ok=True) + generator.eval() + with torch.no_grad(): + samples = generator(noise).cpu() + generator.train() + grid = vutils.make_grid(samples, nrow=grid_size, normalize=True, value_range=(-1, 1)) + timestamp = datetime.utcnow().strftime("%Y%m%dT%H%M%SZ") + out_path = directory_path / f"step_{step:07d}_{timestamp}.png" + vutils.save_image(grid, out_path) + + +def save_checkpoint( + generator: Generator, + discriminator: Discriminator, + optim_g: optim.Optimizer, + optim_d: optim.Optimizer, + directory: str, + epoch: int, + step: int, +) -> None: + directory_path = Path(directory) + directory_path.mkdir(parents=True, exist_ok=True) + torch.save( + { + "epoch": epoch, + "step": step, + "generator": generator.state_dict(), + "discriminator": discriminator.state_dict(), + "optim_g": optim_g.state_dict(), + "optim_d": optim_d.state_dict(), + }, + directory_path / f"dcgan_{step:07d}.pt", + ) + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description="PyTorch DCGAN trainer for CelebA") + parser.add_argument("--data-root", default="./data/img_align_celeba/img_align_celeba", help="Path to the CelebA image directory.") + parser.add_argument("--epochs", type=int, default=50, help="Number of training epochs.") + parser.add_argument("--batch-size", type=int, default=64, help="Training batch size.") + parser.add_argument("--latent-dim", type=int, default=100, help="Dimension of the latent vector z.") + parser.add_argument("--channels", type=int, default=64, help="Base number of convolution channels.") + parser.add_argument("--lr", type=float, default=2e-4, help="Learning rate for both optimizers.") + parser.add_argument("--workers", type=int, default=4, help="Number of dataloader worker processes.") + parser.add_argument("--no-cuda", action="store_true", help="Disable CUDA even if available.") + parser.add_argument("--log-interval", type=int, default=50, help="Steps between logging training metrics.") + parser.add_argument("--sample-interval", type=int, default=500, help="Steps between writing sample grids.") + parser.add_argument("--sample-grid", type=int, default=8, help="Number of images per row/column when saving samples.") + parser.add_argument("--checkpoint-interval", type=int, default=1000, help="Steps between writing checkpoints.") + parser.add_argument("--sample-dir", default="./sample_images_while_training/celeba_dcgan_torch", help="Directory to store generated samples.") + parser.add_argument("--checkpoint-dir", default="./checkpoints/celeba_dcgan_torch", help="Directory to store checkpoints.") + return parser.parse_args() + + +if __name__ == "__main__": + train(parse_args())