下面是基于 GAN 的图像修复 Python 完整代码示例。此代码分为生成器、判别器的构建，以及训练循环的设计，并提供了测试阶段的脚本说明。 ### 1. 导入必要的库 ```python import torch import torch.nn as nn from torchvision import datasets, transforms from torch.utils.data import DataLoader import matplotlib.pyplot as plt import numpy as np import os ``` --- ### 2. 数据预处理与加载 ```python transform = transforms.Compose([ transforms.Resize((128, 128)), transforms.ToTensor(), transforms.Normalize(mean=[0.5], std=[0.5]) ]) dataset = datasets.ImageFolder(root='path_to_dataset', transform=transform) dataloader = DataLoader(dataset, batch_size=16, shuffle=True) ``` --- ### 3. 构建生成器（Generator） 生成器的作用是从噪声中生成新的图像补全缺失区域。 ```python class Generator(nn.Module): def __init__(self): super(Generator, self).__init__() self.main = nn.Sequential( nn.ConvTranspose2d(100, 512, kernel_size=4, stride=1, padding=0, bias=False), nn.BatchNorm2d(512), nn.ReLU(True), nn.ConvTranspose2d(512, 256, kernel_size=4, stride=2, padding=1, bias=False), nn.BatchNorm2d(256), nn.ReLU(True), nn.ConvTranspose2d(256, 128, kernel_size=4, stride=2, padding=1, bias=False), nn.BatchNorm2d(128), nn.ReLU(True), nn.ConvTranspose2d(128, 64, kernel_size=4, stride=2, padding=1, bias=False), nn.BatchNorm2d(64), nn.ReLU(True), nn.ConvTranspose2d(64, 3, kernel_size=4, stride=2, padding=1, bias=False), nn.Tanh() ) def forward(self, input): return self.main(input) ``` --- ### 4. 构建判别器（Discriminator） 判别器的任务是区分真实图像和生成的假图像。 ```python class Discriminator(nn.Module): def __init__(self): super(Discriminator, self).__init__() self.main = nn.Sequential( nn.Conv2d(3, 64, kernel_size=4, stride=2, padding=1, bias=False), nn.LeakyReLU(0.2, inplace=True), nn.Conv2d(64, 128, kernel_size=4, stride=2, padding=1, bias=False), nn.BatchNorm2d(128), nn.LeakyReLU(0.2, inplace=True), nn.Conv2d(128, 256, kernel_size=4, stride=2, padding=1, bias=False), nn.BatchNorm2d(256), nn.LeakyReLU(0.2, inplace=True), nn.Conv2d(256, 512, kernel_size=4, stride=2, padding=1, bias=False), nn.BatchNorm2d(512), nn.LeakyReLU(0.2, inplace=True), nn.Conv2d(512, 1, kernel_size=4, stride=1, padding=0, bias=False), nn.Sigmoid() ) def forward(self, input): return self.main(input).view(-1, 1).squeeze(1) ``` --- ### 5. 初始化模型与超参数 ```python device = torch.device("cuda" if torch.cuda.is_available() else "cpu") netG = Generator().to(device) netD = Discriminator().to(device) criterion = nn.BCELoss() fixed_noise = torch.randn(16, 100, 1, 1, device=device) real_label = 1 fake_label = 0 optimizerD = torch.optim.Adam(netD.parameters(), lr=0.0002, betas=(0.5, 0.999)) optimizerG = torch.optim.Adam(netG.parameters(), lr=0.0002, betas=(0.5, 0.999)) ``` --- ### 6. 训练过程 ```python num_epochs = 50 for epoch in range(num_epochs): for i, data in enumerate(dataloader, 0): netD.zero_grad() real_images = data[0].to(device) b_size = real_images.size(0) labels_real = torch.full((b_size,), real_label, dtype=torch.float, device=device) output_D_real = netD(real_images).view(-1) errD_real = criterion(output_D_real, labels_real) errD_real.backward() noise = torch.randn(b_size, 100, 1, 1, device=device) fake_images = netG(noise) labels_fake = torch.full((b_size,), fake_label, dtype=torch.float, device=device) output_D_fake = netD(fake_images.detach()).view(-1) errD_fake = criterion(output_D_fake, labels_fake) errD_fake.backward() optimizerD.step() netG.zero_grad() output_G = netD(fake_images).view(-1) errG = criterion(output_G, labels_real) errG.backward() optimizerG.step() print(f"[Epoch {epoch}/{num_epochs}] Loss D: {errD_real.item()+errD_fake.item()}, Loss G: {errG.item()}") torch.save(netG.state_dict(), 'generator_model.pth') torch.save(netD.state_dict(), 'discriminator_model.pth') ``` --- ### 7. 测试阶段 测试时可以通过加载已保存的生成器权重文件 `generator_model.pth` 来完成图像修复任务。 ```python def generate_image(model_path, save_dir="output"): netG.load_state_dict(torch.load(model_path)) netG.eval() with torch.no_grad(): generated_images = netG(fixed_noise).detach().cpu() os.makedirs(save_dir, exist_ok=True) for idx, img in enumerate(generated_images): img = (img.permute(1, 2, 0).numpy() * 0.5 + 0.5) * 255 img = img.astype(np.uint8) plt.imsave(os.path.join(save_dir, f'generated_{idx}.png'), img) generate_image('generator_model.pth') ``` --- #### 注意事项 - 上述代码假设输入图像是 RGB 格式的彩色图片[^1]。 - 如果需要对灰度图像进行修复，则应修改通道数为单通道（即 `in_channels=1` 和 `out_channels=1`）[^2]。 - 在实际应用中可能还需要加入额外的正则化项以提高稳定性，例如梯度惩罚或谱归一化等技术[^3]。