PyTorch中GAN每个epoch保存不同单张生成图像的方法

一、生成图像重复的原因与修复 #

你的代码存在两个核心问题导致生成图像不更新：

• 错误遍历数据集而非数据加载器：训练循环里用enumerate(dataset)会直接遍历单张图片，batch_size被强制设为1，模型无法学习到数据分布，Generator参数几乎无更新，导致生成图像完全一致。必须改为遍历loader。
• 数据维度处理疏漏：确保批量图像正确展平、转移到指定设备，避免梯度无法正常传播。

二、保存单张生成图像的修改方法 #

无需使用make_grid拼接网格，直接从固定噪声生成的结果中取出单张图像，调用save_image保存即可。

修正后的完整代码 #

模型与初始化部分 #

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
import torchvision.transforms as transforms
from torchvision.utils import save_image

class Discriminator(nn.Module):
    def __init__(self, img_dim):
        super().__init__()
        self.disc = nn.Sequential(
            nn.Linear(img_dim, 128),
            nn.LeakyReLU(0.1),
            nn.Linear(128, 1),
            nn.Sigmoid()
        )

    def forward(self, x):
        return self.disc(x)
        
class Generator(nn.Module):
    def __init__(self, z_dim, img_dim):
        super().__init__()
        self.gen = nn.Sequential(
            nn.Linear(z_dim, 512),
            nn.LeakyReLU(0.1),
            nn.Linear(512, 1024),
            nn.LeakyReLU(0.1),
            nn.Linear(1024, img_dim),
            nn.Tanh() 
        )

    def forward(self, x):
        return self.gen(x)

# 设备配置
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

lr = 3e-4
z_dim = 64
image_dim = 256 * 256 * 3
batch_size = 32
num_epochs = 16

disc = Discriminator(image_dim).to(device)
gen = Generator(z_dim, image_dim).to(device)
# 生成单张固定噪声，用于观察训练变化
fixed_noise = torch.randn((1, z_dim)).to(device)

# 示例数据预处理（根据你的数据集调整）
transform = transforms.Compose([
    transforms.Resize((256, 256)),
    transforms.ToTensor(),
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])
# dataset = load_dataset(data_path="mountain_dataset", transform=transform)
loader = DataLoader(dataset, batch_size=batch_size, num_workers=0, shuffle=True)
opt_disc = optim.Adam(disc.parameters(), lr=lr)
opt_gen = optim.Adam(gen.parameters(), lr=lr)
criterion = nn.BCELoss()

训练与图像保存部分 #

step = 0

for epoch in range(num_epochs):
    # 修正：遍历DataLoader获取批量数据
    for batch_idx, real in enumerate(loader):
        # 若数据集返回(图像,标签)元组，需改为real[0].to(device)
        real = real.to(device)
        batch_size = real.shape[0]
        real = real.view(-1, image_dim)

        ### 训练判别器
        noise = torch.randn(batch_size, z_dim).to(device)
        fake = gen(noise)
        disc_real = disc(real).view(-1)
        lossD_real = criterion(disc_real, torch.ones_like(disc_real))

        disc_fake = disc(fake).view(-1)
        lossD_fake = criterion(disc_fake, torch.zeros_like(disc_fake))
        lossD = (lossD_real + lossD_fake) / 2

        disc.zero_grad()
        lossD.backward(retain_graph=True)
        opt_disc.step()

        ### 训练生成器
        output = disc(fake).view(-1)
        lossG = criterion(output, torch.ones_like(output))
        gen.zero_grad()
        lossG.backward()
        opt_gen.step()

        if batch_idx == 0:
            print(f"Epoch: [{epoch+1}/{num_epochs}]")

            with torch.no_grad():
                # 生成单张图像并调整维度
                fake_single = gen(fixed_noise).reshape(3, 256, 256)
                # 直接保存单张图像
                save_image(fake_single, f"generated_images/epoch{epoch+1}.png", normalize=True)
                
                step += 1

额外注意事项 #

• 提前创建generated_images文件夹，避免保存图像时路径错误
• 若数据集返回图像与标签的元组，需在训练时提取图像部分（real = real[0].to(device)）
• 可尝试添加BatchNorm层、调整LeakyReLU斜率或学习率，提升GAN训练稳定性

内容的提问来源于stack exchange，提问作者Cihan Yalçın