第 12 章 · 28×28 手写数字 GAN 代码实验¶

# 清晰手写数字 GAN：把低分辨率公开样本放大到 28x28，便于观察生成质量。
torch.manual_seed(7)
torch.set_num_threads(1)
digits = load_digits()
target_digit = 8
digit_images = []
for image, label in zip(digits.images, digits.target):
    if label == target_digit:
        pil_image = Image.fromarray((image / 16.0 * 255).astype("uint8")).resize((28, 28), Image.Resampling.BICUBIC)
        digit_images.append(np.asarray(pil_image).astype("float32") / 255.0)
real_digit_np = np.array(digit_images)
real_digits = torch.tensor(real_digit_np.reshape(len(real_digit_np), -1) * 2 - 1, dtype=torch.float32)

latent_dim = 16
image_dim = real_digits.shape[1]
batch_size = 64

generator = nn.Sequential(
    nn.Linear(latent_dim, 128),
    nn.LeakyReLU(0.2),
    nn.Linear(128, 256),
    nn.LeakyReLU(0.2),
    nn.Linear(256, image_dim),
    nn.Tanh(),
)
discriminator = nn.Sequential(
    nn.Linear(image_dim, 256),
    nn.LeakyReLU(0.2),
    nn.Linear(256, 128),
    nn.LeakyReLU(0.2),
    nn.Linear(128, 1),
)
loss_fn = nn.BCEWithLogitsLoss()
opt_g = torch.optim.Adam(generator.parameters(), lr=0.0005, betas=(0.5, 0.999))
opt_d = torch.optim.Adam(discriminator.parameters(), lr=0.0005, betas=(0.5, 0.999))
gan_rows = []

for step in range(1, 2601):
    idx = torch.randint(0, len(real_digits), (batch_size,))
    real_batch = real_digits[idx]
    z = torch.randn(batch_size, latent_dim)
    fake_batch = generator(z).detach() + 0.03 * torch.randn_like(real_batch)

    real_logits = discriminator(real_batch + 0.03 * torch.randn_like(real_batch))
    fake_logits = discriminator(fake_batch)
    d_loss = loss_fn(real_logits, torch.full_like(real_logits, 0.9)) + loss_fn(fake_logits, torch.full_like(fake_logits, 0.1))
    opt_d.zero_grad()
    d_loss.backward()
    opt_d.step()

    z = torch.randn(batch_size, latent_dim)
    generated = generator(z)
    g_logits = discriminator(generated)
    anchor = real_digits[torch.randint(0, len(real_digits), (batch_size,))]
    g_loss = loss_fn(g_logits, torch.ones_like(g_logits)) + 0.04 * ((generated - anchor) ** 2).mean()
    opt_g.zero_grad()
    g_loss.backward()
    opt_g.step()

    if step % 400 == 0:
        with torch.no_grad():
            gan_rows.append({
                "训练步": step,
                "判别器损失": float(d_loss),
                "生成器损失": float(g_loss),
                "真实样本评分": float(torch.sigmoid(discriminator(real_batch)).mean()),
                "生成样本评分": float(torch.sigmoid(discriminator(generator(torch.randn(batch_size, latent_dim)))).mean()),
            })

with torch.no_grad():
    z_fixed = torch.randn(16, latent_dim)
    gan_samples = generator(z_fixed).reshape(16, 28, 28).numpy()
gan_trace = pd.DataFrame(gan_rows)
display(pd.DataFrame({
    "数据集": ["sklearn 手写数字"],
    "目标数字": [target_digit],
    "真实样本数": [len(real_digits)],
    "图像尺寸": ["28x28"],
    "生成向量维度": [image_dim],
}))
display(gan_trace.round(3))

# 绘制训练曲线、真实数字和生成数字。
fig = plt.figure(figsize=(10.8, 7.2))
gs = fig.add_gridspec(2, 4, height_ratios=[1.0, 1.6], hspace=0.38, wspace=0.18)
ax_loss = fig.add_subplot(gs[0, :2])
ax_score = fig.add_subplot(gs[0, 2:])
ax_loss.plot(gan_trace["训练步"], gan_trace["判别器损失"], color="#2563eb", linewidth=2.0, label="判别器损失")
ax_loss.plot(gan_trace["训练步"], gan_trace["生成器损失"], color="#f97316", linewidth=2.0, label="生成器损失")
ax_loss.set_title("训练损失", loc="left", fontweight="bold")
ax_loss.grid(True, color="#e2e8f0", linewidth=0.8)
ax_loss.legend()

ax_score.plot(gan_trace["训练步"], gan_trace["真实样本评分"], color="#16a34a", linewidth=2.0, label="真实样本评分")
ax_score.plot(gan_trace["训练步"], gan_trace["生成样本评分"], color="#dc2626", linewidth=2.0, label="生成样本评分")
ax_score.set_title("判别器输出", loc="left", fontweight="bold")
ax_score.grid(True, color="#e2e8f0", linewidth=0.8)
ax_score.legend()

def digit_tile(images, grid=4, pad=4):
    h, w = images[0].shape
    canvas = np.ones((grid * h + (grid - 1) * pad, grid * w + (grid - 1) * pad), dtype=float)
    for i in range(grid):
        for j in range(grid):
            r0 = i * (h + pad)
            c0 = j * (w + pad)
            canvas[r0:r0 + h, c0:c0 + w] = images[i * grid + j]
    return canvas

real_ids = np.linspace(0, len(real_digit_np) - 1, 16, dtype=int)
real_tile = digit_tile(real_digit_np[real_ids])
fake_tile = digit_tile(((gan_samples[:16] + 1) / 2).clip(0, 1))
ax_real = fig.add_subplot(gs[1, :2])
ax_fake = fig.add_subplot(gs[1, 2:])
ax_real.imshow(real_tile, cmap="gray_r", vmin=0, vmax=1)
ax_real.set_title("真实数字 8", loc="left", fontweight="bold")
ax_fake.imshow(fake_tile, cmap="gray_r", vmin=0, vmax=1)
ax_fake.set_title("生成数字 8", loc="left", fontweight="bold")
for ax in (ax_real, ax_fake):
    ax.set_xticks([])
    ax.set_yticks([])
fig.suptitle("GAN：生成器试图让判别器把生成数字当成真实数字", x=0.08, ha="left", fontsize=14, fontweight="bold", color="#0f172a")
plt.tight_layout()
plt.show()