【计算机视觉技术 - 人脸生成】2.GAN网络的构建和训练

         GAN 是一种常用的优秀的图像生成模型。我们使用了支持条件生成的 cGAN。下面介绍简单 cGAN 模型的构建以及训练过程。

2.1 在 model 文件夹中新建 nets.py 文件

import torch
import torch.nn as nn# 生成器类
class Generator(nn.Module):def __init__(self, nz=100, nc=3, ngf=128, num_classes=4):super(Generator, self).__init__()self.label_emb = nn.Embedding(num_classes, nz)self.main = nn.Sequential(nn.ConvTranspose2d(nz + nz, ngf * 8, 4, 1, 0, bias=False),nn.BatchNorm2d(ngf * 8),nn.ReLU(True),nn.ConvTranspose2d(ngf * 8, ngf * 4, 4, 2, 1, bias=False),nn.BatchNorm2d(ngf * 4),nn.ReLU(True),nn.ConvTranspose2d(ngf * 4, ngf * 2, 4, 2, 1, bias=False),nn.BatchNorm2d(ngf * 2),nn.ReLU(True),nn.ConvTranspose2d(ngf * 2, ngf, 4, 2, 1, bias=False),nn.BatchNorm2d(ngf),nn.ReLU(True),nn.ConvTranspose2d(ngf, nc, 4, 2, 1, bias=False),nn.Tanh())def forward(self, z, labels):c = self.label_emb(labels).unsqueeze(2).unsqueeze(3)x = torch.cat([z, c], 1)return self.main(x)# 判别器类
class Discriminator(nn.Module):def __init__(self, nc=3, ndf=64, num_classes=4):super(Discriminator, self).__init__()self.label_emb = nn.Embedding(num_classes, nc * 64 * 64)self.main = nn.Sequential(nn.Conv2d(nc + 1, ndf, 4, 2, 1, bias=False),nn.LeakyReLU(0.2, inplace=True),nn.Conv2d(ndf, ndf * 2, 4, 2, 1, bias=False),nn.BatchNorm2d(ndf * 2),nn.LeakyReLU(0.2, inplace=True),nn.Conv2d(ndf * 2, ndf * 4, 4, 2, 1, bias=False),nn.BatchNorm2d(ndf * 4),nn.LeakyReLU(0.2, inplace=True),nn.Conv2d(ndf * 4, 1, 4, 1, 0, bias=False),nn.Sigmoid())def forward(self, img, labels):c = self.label_emb(labels).view(labels.size(0), 1, 64, 64)x = torch.cat([img, c], 1)return self.main(x)

2.2新建cGAN_net.py

import torch
import torch.nn as nn
from torch.optim import Adam
from torchvision import datasets, transforms
from torch.utils.data import DataLoader
from torch.optim.lr_scheduler import StepLR# ===========================
# Conditional DCGAN 实现
# ===========================
class cDCGAN:def __init__(self, data_root, batch_size, device, latent_dim=100, num_classes=4):self.device = deviceself.batch_size = batch_sizeself.latent_dim = latent_dimself.num_classes = num_classes# 数据加载器self.train_loader = self.get_dataloader(data_root)# 初始化生成器和判别器self.generator = self.build_generator().to(device)self.discriminator = self.build_discriminator().to(device)# 初始化权重self.generator.apply(self.weights_init)self.discriminator.apply(self.weights_init)# 损失函数和优化器self.criterion = nn.BCELoss()self.optimizer_G = Adam(self.generator.parameters(), lr=0.0001, betas=(0.5, 0.999))self.optimizer_D = Adam(self.discriminator.parameters(), lr=0.0001, betas=(0.5, 0.999))# 学习率调度器self.scheduler_G = StepLR(self.optimizer_G, step_size=10, gamma=0.5)  # 每10个epoch学习率减半self.scheduler_D = StepLR(self.optimizer_D, step_size=10, gamma=0.5)def get_dataloader(self, data_root):transform = transforms.Compose([transforms.Resize(128),transforms.ToTensor(),transforms.Normalize((0.5,), (0.5,))])dataset = datasets.ImageFolder(root=data_root, transform=transform)return DataLoader(dataset, batch_size=self.batch_size, shuffle=True,num_workers=8, pin_memory=True, persistent_workers=True)@staticmethoddef weights_init(model):"""权重初始化"""if isinstance(model, (nn.Conv2d, nn.Linear)):nn.init.normal_(model.weight.data, 0.0, 0.02)if model.bias is not None:nn.init.constant_(model.bias.data, 0)def train_step(self, epoch, step, num_epochs):"""单次训练步骤"""self.generator.train()self.discriminator.train()G_losses, D_losses = [], []for i, (real_img, labels) in enumerate(self.train_loader):# 确保 real_img 和 labels 在同一设备real_img = real_img.to(self.device)labels = labels.to(self.device)batch_size = real_img.size(0)# # 标签   11.19 15:11:12修改# valid = torch.ones((batch_size, 1), device=self.device)# fake = torch.zeros((batch_size, 1), device=self.device)# 标签平滑# smooth_valid = torch.full((batch_size, 1), 1, device=self.device)  # 平滑真实标签# smooth_fake = torch.full((batch_size, 1), 0, device=self.device)  # 平滑伪造标签# smooth_valid = torch.full((batch_size, 1), torch.rand(1).item() * 0.1 + 0.9, device=self.device)# smooth_fake = torch.full((batch_size, 1), torch.rand(1).item() * 0.1, device=self.device)# smooth_valid = torch.full((batch_size, 1), max(0.7, 1 - epoch * 0.001), device=self.device)# smooth_fake = torch.full((batch_size, 1), min(0.3, epoch * 0.001), device=self.device)# 动态调整标签范围smooth_valid = torch.full((batch_size, 1), max(0.9, 1 - 0.0001 * epoch), device=self.device)smooth_fake = torch.full((batch_size, 1), min(0.1, 0.0001 * epoch), device=self.device)# 替换以下两处代码valid = smooth_validfake = smooth_fake# ========== 训练判别器 ==========real_pred = self.discriminator(real_img, labels)# d_real_loss = self.criterion(real_pred, valid)d_real_loss = self.criterion(real_pred, valid - 0.1 * torch.rand_like(valid))noise = torch.randn(batch_size, self.latent_dim, device=self.device)# gen_labels = torch.randint(0, self.num_classes, (batch_size,), device=self.device)gen_labels = torch.randint(0, self.num_classes, (batch_size,), device=self.device) + torch.randint(-1, 2, (batch_size,), device=self.device)gen_labels = torch.clamp(gen_labels, 0, self.num_classes - 1)  # 确保标签在范围内gen_img = self.generator(noise, gen_labels)fake_pred = self.discriminator(gen_img.detach(), gen_labels)# d_fake_loss = self.criterion(fake_pred, fake)d_fake_loss = self.criterion(fake_pred, fake + 0.1 * torch.rand_like(fake))d_loss = (d_real_loss + d_fake_loss) / 2self.optimizer_D.zero_grad()d_loss.backward()self.optimizer_D.step()D_losses.append(d_loss.item())# ========== 训练生成器 ==========gen_pred = self.discriminator(gen_img, gen_labels)g_loss = self.criterion(gen_pred, valid)self.optimizer_G.zero_grad()g_loss.backward()self.optimizer_G.step()G_losses.append(g_loss.item())print(f'第 {epoch}/{num_epochs} 轮, Batch {i + 1}/{len(self.train_loader)}, 'f'D Loss: {d_loss:.4f}, G Loss: {g_loss:.4f}')step += 1return G_losses, D_losses, stepdef build_generator(self):"""生成器"""return Generator(latent_dim=self.latent_dim, num_classes=self.num_classes)def build_discriminator(self):"""判别器"""return Discriminator(num_classes=self.num_classes)def load_model(self, model_path):"""加载模型权重"""checkpoint = torch.load(model_path, map_location=self.device)self.generator.load_state_dict(checkpoint['generator_state_dict'])self.optimizer_G.load_state_dict(checkpoint['optimizer_G_state_dict'])self.discriminator.load_state_dict(checkpoint['discriminator_state_dict'])self.optimizer_D.load_state_dict(checkpoint['optimizer_D_state_dict'])epoch = checkpoint['epoch']print(f"加载了模型权重，起始训练轮次为 {epoch}")return epochdef save_model(self, epoch, save_path):"""保存模型"""torch.save({'epoch': epoch,'scheduler_G_state_dict': self.scheduler_G.state_dict(),'scheduler_D_state_dict': self.scheduler_D.state_dict(),'generator_state_dict': self.generator.state_dict(),'optimizer_G_state_dict': self.optimizer_G.state_dict(),'discriminator_state_dict': self.discriminator.state_dict(),'optimizer_D_state_dict': self.optimizer_D.state_dict(),}, save_path)print(f"模型已保存至 {save_path}")# ===========================
# 生成器
# ===========================
class Generator(nn.Module):def __init__(self, latent_dim=100, num_classes=4, img_channels=3):super(Generator, self).__init__()self.latent_dim = latent_dimself.label_emb = nn.Embedding(num_classes, num_classes)self.init_size = 8self.l1 = nn.Linear(latent_dim + num_classes, 256 * self.init_size * self.init_size)self.conv_blocks = nn.Sequential(nn.BatchNorm2d(256),nn.Upsample(scale_factor=2),nn.Conv2d(256, 128, 3, padding=1),nn.BatchNorm2d(128),nn.LeakyReLU(0.2, inplace=True),nn.Upsample(scale_factor=2),nn.Conv2d(128, 64, 3, padding=1),nn.BatchNorm2d(64),nn.LeakyReLU(0.2, inplace=True),nn.Upsample(scale_factor=2),nn.Conv2d(64, 32, 3, padding=1),nn.BatchNorm2d(32),nn.LeakyReLU(0.2, inplace=True),nn.Upsample(scale_factor=2),nn.Conv2d(32, img_channels, 3, padding=1),nn.Tanh())def forward(self, noise, labels):labels = labels.to(self.label_emb.weight.device)label_embedding = self.label_emb(labels)x = torch.cat((noise, label_embedding), dim=1)x = self.l1(x).view(x.size(0), 256, self.init_size, self.init_size)return self.conv_blocks(x)# ===========================
# 判别器
# ===========================
class Discriminator(nn.Module):def __init__(self, img_channels=3, num_classes=4):super(Discriminator, self).__init__()self.label_embedding = nn.Embedding(num_classes, img_channels)self.model = nn.Sequential(nn.Conv2d(img_channels * 2, 64, 4, stride=2, padding=1),nn.LeakyReLU(0.2, inplace=True),nn.Conv2d(64, 128, 4, stride=2, padding=1),nn.BatchNorm2d(128),nn.LeakyReLU(0.2, inplace=True),nn.Conv2d(128, 256, 4, stride=2, padding=1),nn.BatchNorm2d(256),nn.LeakyReLU(0.2, inplace=True),nn.Conv2d(256, 512, 4, stride=2, padding=1),nn.BatchNorm2d(512),nn.LeakyReLU(0.2, inplace=True))self.output_layer = nn.Sequential(nn.Linear(512 * 8 * 8, 1),nn.Sigmoid())def forward(self, img, labels):labels = labels.to(self.label_embedding.weight.device)label_embedding = self.label_embedding(labels).unsqueeze(2).unsqueeze(3)label_embedding = label_embedding.expand(-1, -1, img.size(2), img.size(3))x = torch.cat((img, label_embedding), dim=1)x = self.model(x).view(x.size(0), -1)return self.output_layer(x)

2.3新建cGAN_trainer.py

import os
import torch
import argparse
from cGAN_net import cDCGAN
from utils import plot_loss, plot_result
import timeos.environ['OMP_NUM_THREADS'] = '1'def main(args):# 初始化设备和训练参数device = torch.device(args.device if torch.cuda.is_available() else 'cpu')model = cDCGAN(data_root=args.data_root, batch_size=args.batch_size, device=device, latent_dim=args.latent_dim)# 添加学习率调度器scheduler_G = torch.optim.lr_scheduler.StepLR(model.optimizer_G, step_size=10, gamma=0.5)scheduler_D = torch.optim.lr_scheduler.StepLR(model.optimizer_D, step_size=10, gamma=0.5)start_epoch = 0# 如果有保存的模型，加载if args.load_model and os.path.exists(args.load_model):start_epoch = model.load_model(args.load_model) + 1# 恢复调度器状态scheduler_G_path = f"{args.load_model}_scheduler_G.pt"scheduler_D_path = f"{args.load_model}_scheduler_D.pt"if os.path.exists(scheduler_G_path) and os.path.exists(scheduler_D_path):scheduler_G.load_state_dict(torch.load(scheduler_G_path))scheduler_D.load_state_dict(torch.load(scheduler_D_path))print(f"成功恢复调度器状态：{scheduler_G_path}, {scheduler_D_path}")else:print("未找到调度器状态文件，使用默认调度器设置")print(f"从第 {start_epoch} 轮继续训练...")print(f"开始训练，从第 {start_epoch + 1} 轮开始...")# 创建保存路径os.makedirs(args.save_dir, exist_ok=True)os.makedirs(os.path.join(args.save_dir, 'log'), exist_ok=True)# 训练循环D_avg_losses, G_avg_losses = [], []for epoch in range(start_epoch, args.epochs):G_losses, D_losses, step = model.train_step(epoch, step=0, num_epochs=args.epochs)# 计算平均损失D_avg_loss = sum(D_losses) / len(D_losses) if D_losses else 0.0G_avg_loss = sum(G_losses) / len(G_losses) if G_losses else 0.0D_avg_losses.append(D_avg_loss)G_avg_losses.append(G_avg_loss)# 保存损失曲线图plot_loss(start_epoch, args.epochs, D_avg_losses, G_avg_losses, epoch + 1, save=True,save_dir=os.path.join(args.save_dir, "log"))# 生成并保存图片labels = torch.tensor([0, 1, 2, 3]).to(device)if (epoch + 1) % args.save_freq == 0:  # 每隔一定轮次保存生成结果z = torch.randn(len(labels), args.latent_dim, device=device)  # 随机生成噪声plot_result(model.generator, z, labels, epoch + 1, save_dir=os.path.join(args.save_dir, 'log'))# 每10个epoch保存模型if (epoch + 1) % args.save_interval == 0:timestamp = int(time.time())save_path = os.path.join(args.save_dir, f"cgan_epoch_{epoch + 1}_{timestamp}.pth")model.save_model(epoch + 1, save_path)print(f"第 {epoch + 1} 轮的模型已保存，保存路径为 {save_path}")# 更新学习率调度器scheduler_G.step()scheduler_D.step()if __name__ == "__main__":parser = argparse.ArgumentParser()parser.add_argument('--data_root', type=str, default='data/crop128', help="数据集根目录")parser.add_argument('--save_dir', type=str, default='./chkpt/cgan_model', help="保存模型的目录")parser.add_argument('--load_model', type=str, default=None, help="要加载的模型路径（可选）")parser.add_argument('--epochs', type=int, default=1000, help="训练的轮数")parser.add_argument('--save_interval', type=int, default=10, help="保存模型检查点的间隔（按轮数）")parser.add_argument('--batch_size', type=int, default=64, help="训练的批次大小")parser.add_argument('--device', type=str, default='cuda', help="使用的设备（如 cuda 或 cpu）")parser.add_argument('--latent_dim', type=int, default=100, help="生成器的潜在空间维度")parser.add_argument('--save_freq', type=int, default=1, help="每隔多少轮保存一次生成结果（默认: 1）")args = parser.parse_args()main(args)

结果分析：

2.4中间结果可视化处理

新建utils.py，编写绘制中间结果和中间损失线图的函数，代码如下：

import matplotlib.pyplot as plt
import numpy as np
import os
import torchdef denorm(x):out = (x + 1) / 2return out.clamp(0, 1)def plot_loss(start_epoch, num_epochs, d_losses, g_losses, num_epoch, save=False, save_dir='celebA_cDCGAN_results/', show=False):"""绘制损失函数曲线，从 start_epoch 到 num_epochs。Args:start_epoch: 起始轮次num_epochs: 总轮次d_losses: 判别器损失列表g_losses: 生成器损失列表num_epoch: 当前训练轮次save: 是否保存绘图save_dir: 保存路径show: 是否显示绘图"""fig, ax = plt.subplots()ax.set_xlim(start_epoch, num_epochs)ax.set_ylim(0, max(np.max(g_losses), np.max(d_losses)) * 1.1)plt.xlabel(f'Epoch {num_epoch + 1}')plt.ylabel('Loss values')plt.plot(d_losses, label='Discriminator')plt.plot(g_losses, label='Generator')plt.legend()if save:if not os.path.exists(save_dir):os.makedirs(save_dir)save_fn = os.path.join(save_dir, f'cDCGAN_losses_epoch.png')plt.savefig(save_fn)if show:plt.show()else:plt.close()def plot_result(generator, z, labels, epoch, save_dir=None, show=False):"""生成并保存或显示生成的图片结果。Args:generator: 生成器模型z: 随机噪声张量labels: 标签张量epoch: 当前训练轮数save_dir: 保存图片的路径（可选）show: 是否显示生成的图片（可选）"""# 调用生成器，生成图像generator.eval()  # 设置为评估模式with torch.no_grad():gen_images = generator(z, labels)  # 同时传入 z 和 labelsgenerator.train()  # 恢复训练模式# 图像反归一化gen_images = denorm(gen_images)# 绘制图片fig, ax = plt.subplots(1, len(gen_images), figsize=(15, 15))for i in range(len(gen_images)):ax[i].imshow(gen_images[i].permute(1, 2, 0).cpu().numpy())  # 转换为可显示格式ax[i].axis('off')# 保存或显示图片if save_dir:os.makedirs(save_dir, exist_ok=True)save_path = os.path.join(save_dir, f'epoch_{epoch}.png')plt.savefig(save_path)if show:plt.show()plt.close(fig)

执行 cGAN_trainer.py 文件，完成模型训练。