SNGAN (Spectral Normalization GAN)の概要とアルゴリズム及び実装例

pip install torch torchvision matplotlib

import torch
import torch.nn as nn
import torch.optim as optim
import torchvision
import torchvision.transforms as transforms
import matplotlib.pyplot as plt

# ===========================
#  1. ハイパーパラメータ設定
# ===========================
latent_dim = 100  # 潜在変数の次元
image_size = 64   # 画像サイズ (64x64)
batch_size = 128
num_epochs = 20
lr = 0.0002
beta1 = 0.5  # Adam optimizerのハイパーパラメータ

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# ===========================
#  2. データセットの準備 (CIFAR-10)
# ===========================
transform = transforms.Compose([
    transforms.Resize(image_size),
    transforms.ToTensor(),
    transforms.Normalize([0.5], [0.5])  # [-1, 1] に正規化
])

dataset = torchvision.datasets.CIFAR10(root="./data", download=True, transform=transform)
dataloader = torch.utils.data.DataLoader(dataset, batch_size=batch_size, shuffle=True)

# ===========================
#  3. 生成器（Generator）の定義
# ===========================
class Generator(nn.Module):
    def __init__(self):
        super(Generator, self).__init__()
        self.model = nn.Sequential(
            nn.ConvTranspose2d(latent_dim, 512, 4, 1, 0, bias=False),
            nn.BatchNorm2d(512),
            nn.ReLU(True),

            nn.ConvTranspose2d(512, 256, 4, 2, 1, bias=False),
            nn.BatchNorm2d(256),
            nn.ReLU(True),

            nn.ConvTranspose2d(256, 128, 4, 2, 1, bias=False),
            nn.BatchNorm2d(128),
            nn.ReLU(True),

            nn.ConvTranspose2d(128, 64, 4, 2, 1, bias=False),
            nn.BatchNorm2d(64),
            nn.ReLU(True),

            nn.ConvTranspose2d(64, 3, 4, 2, 1, bias=False),
            nn.Tanh()  # 出力を [-1, 1] にする
        )

    def forward(self, z):
        return self.model(z)

# ===========================
#  4. 識別器（Discriminator）の定義 (Spectral Normalization)
# ===========================
class Discriminator(nn.Module):
    def __init__(self):
        super(Discriminator, self).__init__()
        self.model = nn.Sequential(
            nn.utils.spectral_norm(nn.Conv2d(3, 64, 4, 2, 1, bias=False)),
            nn.LeakyReLU(0.2, inplace=True),

            nn.utils.spectral_norm(nn.Conv2d(64, 128, 4, 2, 1, bias=False)),
            nn.BatchNorm2d(128),
            nn.LeakyReLU(0.2, inplace=True),

            nn.utils.spectral_norm(nn.Conv2d(128, 256, 4, 2, 1, bias=False)),
            nn.BatchNorm2d(256),
            nn.LeakyReLU(0.2, inplace=True),

            nn.utils.spectral_norm(nn.Conv2d(256, 512, 4, 2, 1, bias=False)),
            nn.BatchNorm2d(512),
            nn.LeakyReLU(0.2, inplace=True),

            nn.Conv2d(512, 1, 4, 1, 0, bias=False)
        )

    def forward(self, x):
        return self.model(x).view(-1, 1).squeeze(1)  # (batch, 1) → (batch,)

# ===========================
#  5. モデルとオプティマイザの設定
# ===========================
netG = Generator().to(device)
netD = Discriminator().to(device)

criterion = nn.BCEWithLogitsLoss()  # BCE loss + logit出力

optimizerG = optim.Adam(netG.parameters(), lr=lr, betas=(beta1, 0.999))
optimizerD = optim.Adam(netD.parameters(), lr=lr, betas=(beta1, 0.999))

# ===========================
#  6. 学習ループ
# ===========================
fixed_noise = torch.randn(16, latent_dim, 1, 1, device=device)

for epoch in range(num_epochs):
    for i, (real_images, _) in enumerate(dataloader):
        real_images = real_images.to(device)
        batch_size = real_images.size(0)

        # ラベル作成
        real_labels = torch.ones(batch_size, device=device)
        fake_labels = torch.zeros(batch_size, device=device)

        # === 識別器（D）の学習 ===
        optimizerD.zero_grad()
        outputs = netD(real_images)
        loss_real = criterion(outputs, real_labels)
        loss_real.backward()

        noise = torch.randn(batch_size, latent_dim, 1, 1, device=device)
        fake_images = netG(noise)
        outputs = netD(fake_images.detach())  # Gの勾配を止める
        loss_fake = criterion(outputs, fake_labels)
        loss_fake.backward()

        optimizerD.step()

        # === 生成器（G）の学習 ===
        optimizerG.zero_grad()
        outputs = netD(fake_images)
        loss_G = criterion(outputs, real_labels)
        loss_G.backward()
        optimizerG.step()

        # 途中経過の出力
        if i % 500 == 0:
            print(f"Epoch [{epoch+1}/{num_epochs}], Step [{i}/{len(dataloader)}], "
                  f"D Loss: {loss_real.item() + loss_fake.item():.4f}, G Loss: {loss_G.item():.4f}")

    # 画像の生成と表示
    with torch.no_grad():
        fake_images = netG(fixed_noise).cpu()
    grid = torchvision.utils.make_grid(fake_images, normalize=True)
    plt.figure(figsize=(6,6))
    plt.imshow(grid.permute(1, 2, 0))
    plt.title(f"Epoch {epoch+1}")
    plt.axis("off")
    plt.show()