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

         GAN 是一种常用的优秀的图像天生模子。我们使用了支持条件天生的 cGAN。下面先容简单 cGAN 模子的构建以及练习过程。
2.1 在 model 文件夹中新建 nets.py 文件

1. import torch
2. import torch.nn as nn
5. # 生成器类
6. class Generator(nn.Module):
7.     def __init__(self, nz=100, nc=3, ngf=128, num_classes=4):
8.         super(Generator, self).__init__()
9.         self.label_emb = nn.Embedding(num_classes, nz)
11.         self.main = nn.Sequential(
12.             nn.ConvTranspose2d(nz + nz, ngf * 8, 4, 1, 0, bias=False),
13.             nn.BatchNorm2d(ngf * 8),
14.             nn.ReLU(True),
16.             nn.ConvTranspose2d(ngf * 8, ngf * 4, 4, 2, 1, bias=False),
17.             nn.BatchNorm2d(ngf * 4),
18.             nn.ReLU(True),
20.             nn.ConvTranspose2d(ngf * 4, ngf * 2, 4, 2, 1, bias=False),
21.             nn.BatchNorm2d(ngf * 2),
22.             nn.ReLU(True),
24.             nn.ConvTranspose2d(ngf * 2, ngf, 4, 2, 1, bias=False),
25.             nn.BatchNorm2d(ngf),
26.             nn.ReLU(True),
28.             nn.ConvTranspose2d(ngf, nc, 4, 2, 1, bias=False),
29.             nn.Tanh()
30.         )
32.     def forward(self, z, labels):
33.         c = self.label_emb(labels).unsqueeze(2).unsqueeze(3)
34.         x = torch.cat([z, c], 1)
35.         return self.main(x)
38. # 判别器类
39. class Discriminator(nn.Module):
40.     def __init__(self, nc=3, ndf=64, num_classes=4):
41.         super(Discriminator, self).__init__()
42.         self.label_emb = nn.Embedding(num_classes, nc * 64 * 64)
44.         self.main = nn.Sequential(
45.             nn.Conv2d(nc + 1, ndf, 4, 2, 1, bias=False),
46.             nn.LeakyReLU(0.2, inplace=True),
48.             nn.Conv2d(ndf, ndf * 2, 4, 2, 1, bias=False),
49.             nn.BatchNorm2d(ndf * 2),
50.             nn.LeakyReLU(0.2, inplace=True),
52.             nn.Conv2d(ndf * 2, ndf * 4, 4, 2, 1, bias=False),
53.             nn.BatchNorm2d(ndf * 4),
54.             nn.LeakyReLU(0.2, inplace=True),
56.             nn.Conv2d(ndf * 4, 1, 4, 1, 0, bias=False),
57.             nn.Sigmoid()
58.         )
60.     def forward(self, img, labels):
61.         c = self.label_emb(labels).view(labels.size(0), 1, 64, 64)
62.         x = torch.cat([img, c], 1)
63.         return self.main(x)

复制代码

2.2新建cGAN_net.py

1. import torch
2. import torch.nn as nn
3. from torch.optim import Adam
4. from torchvision import datasets, transforms
5. from torch.utils.data import DataLoader
6. from torch.optim.lr_scheduler import StepLR
9. # ===========================
10. # Conditional DCGAN 实现
11. # ===========================
12. class cDCGAN:
13.     def __init__(self, data_root, batch_size, device, latent_dim=100, num_classes=4):
14.         self.device = device
15.         self.batch_size = batch_size
16.         self.latent_dim = latent_dim
17.         self.num_classes = num_classes
19.         # 数据加载器
20.         self.train_loader = self.get_dataloader(data_root)
22.         # 初始化生成器和判别器
23.         self.generator = self.build_generator().to(device)
24.         self.discriminator = self.build_discriminator().to(device)
26.         # 初始化权重
27.         self.generator.apply(self.weights_init)
28.         self.discriminator.apply(self.weights_init)
30.         # 损失函数和优化器
31.         self.criterion = nn.BCELoss()
32.         self.optimizer_G = Adam(self.generator.parameters(), lr=0.0001, betas=(0.5, 0.999))
33.         self.optimizer_D = Adam(self.discriminator.parameters(), lr=0.0001, betas=(0.5, 0.999))
35.         # 学习率调度器
36.         self.scheduler_G = StepLR(self.optimizer_G, step_size=10, gamma=0.5)  # 每10个epoch学习率减半
37.         self.scheduler_D = StepLR(self.optimizer_D, step_size=10, gamma=0.5)
39.     def get_dataloader(self, data_root):
40.         transform = transforms.Compose([
41.             transforms.Resize(128),
42.             transforms.ToTensor(),
43.             transforms.Normalize((0.5,), (0.5,))
44.         ])
45.         dataset = datasets.ImageFolder(root=data_root, transform=transform)
46.         return DataLoader(dataset, batch_size=self.batch_size, shuffle=True,
47.                           num_workers=8, pin_memory=True, persistent_workers=True)
49.     @staticmethod
50.     def weights_init(model):
51.         """权重初始化"""
52.         if isinstance(model, (nn.Conv2d, nn.Linear)):
53.             nn.init.normal_(model.weight.data, 0.0, 0.02)
54.             if model.bias is not None:
55.                 nn.init.constant_(model.bias.data, 0)
57.     def train_step(self, epoch, step, num_epochs):
58.         """单次训练步骤"""
59.         self.generator.train()
60.         self.discriminator.train()
62.         G_losses, D_losses = [], []
63.         for i, (real_img, labels) in enumerate(self.train_loader):
64.             # 确保 real_img 和 labels 在同一设备
65.             real_img = real_img.to(self.device)
66.             labels = labels.to(self.device)
68.             batch_size = real_img.size(0)
70.             # # 标签   11.19 15:11:12修改
71.             # valid = torch.ones((batch_size, 1), device=self.device)
72.             # fake = torch.zeros((batch_size, 1), device=self.device)
74.             # 标签平滑
75.             # smooth_valid = torch.full((batch_size, 1), 1, device=self.device)  # 平滑真实标签
76.             # smooth_fake = torch.full((batch_size, 1), 0, device=self.device)  # 平滑伪造标签
77.             # smooth_valid = torch.full((batch_size, 1), torch.rand(1).item() * 0.1 + 0.9, device=self.device)
78.             # smooth_fake = torch.full((batch_size, 1), torch.rand(1).item() * 0.1, device=self.device)
79.             # smooth_valid = torch.full((batch_size, 1), max(0.7, 1 - epoch * 0.001), device=self.device)
80.             # smooth_fake = torch.full((batch_size, 1), min(0.3, epoch * 0.001), device=self.device)
81.             # 动态调整标签范围
82.             smooth_valid = torch.full((batch_size, 1), max(0.9, 1 - 0.0001 * epoch), device=self.device)
83.             smooth_fake = torch.full((batch_size, 1), min(0.1, 0.0001 * epoch), device=self.device)
85.             # 替换以下两处代码
86.             valid = smooth_valid
87.             fake = smooth_fake
89.             # ========== 训练判别器 ==========
90.             real_pred = self.discriminator(real_img, labels)
91.             # d_real_loss = self.criterion(real_pred, valid)
92.             d_real_loss = self.criterion(real_pred, valid - 0.1 * torch.rand_like(valid))
94.             noise = torch.randn(batch_size, self.latent_dim, device=self.device)
95.             # gen_labels = torch.randint(0, self.num_classes, (batch_size,), device=self.device)
96.             gen_labels = torch.randint(0, self.num_classes, (batch_size,), device=self.device) + torch.randint(-1, 2, (
97.             batch_size,), device=self.device)
98.             gen_labels = torch.clamp(gen_labels, 0, self.num_classes - 1)  # 确保标签在范围内
100.             gen_img = self.generator(noise, gen_labels)
102.             fake_pred = self.discriminator(gen_img.detach(), gen_labels)
103.             # d_fake_loss = self.criterion(fake_pred, fake)
104.             d_fake_loss = self.criterion(fake_pred, fake + 0.1 * torch.rand_like(fake))
106.             d_loss = (d_real_loss + d_fake_loss) / 2
107.             self.optimizer_D.zero_grad()
108.             d_loss.backward()
109.             self.optimizer_D.step()
110.             D_losses.append(d_loss.item())
112.             # ========== 训练生成器 ==========
113.             gen_pred = self.discriminator(gen_img, gen_labels)
114.             g_loss = self.criterion(gen_pred, valid)
116.             self.optimizer_G.zero_grad()
117.             g_loss.backward()
118.             self.optimizer_G.step()
119.             G_losses.append(g_loss.item())
121.             print(f'第 {epoch}/{num_epochs} 轮, Batch {i + 1}/{len(self.train_loader)}, '
122.                   f'D Loss: {d_loss:.4f}, G Loss: {g_loss:.4f}')
124.         step += 1
125.         return G_losses, D_losses, step
127.     def build_generator(self):
128.         """生成器"""
129.         return Generator(latent_dim=self.latent_dim, num_classes=self.num_classes)
131.     def build_discriminator(self):
132.         """判别器"""
133.         return Discriminator(num_classes=self.num_classes)
135.     def load_model(self, model_path):
136.         """加载模型权重"""
137.         checkpoint = torch.load(model_path, map_location=self.device)
138.         self.generator.load_state_dict(checkpoint['generator_state_dict'])
139.         self.optimizer_G.load_state_dict(checkpoint['optimizer_G_state_dict'])
140.         self.discriminator.load_state_dict(checkpoint['discriminator_state_dict'])
141.         self.optimizer_D.load_state_dict(checkpoint['optimizer_D_state_dict'])
142.         epoch = checkpoint['epoch']
143.         print(f"加载了模型权重，起始训练轮次为 {epoch}")
144.         return epoch
146.     def save_model(self, epoch, save_path):
147.         """保存模型"""
148.         torch.save({
149.             'epoch': epoch,
150.             'scheduler_G_state_dict': self.scheduler_G.state_dict(),
151.             'scheduler_D_state_dict': self.scheduler_D.state_dict(),
152.             'generator_state_dict': self.generator.state_dict(),
153.             'optimizer_G_state_dict': self.optimizer_G.state_dict(),
154.             'discriminator_state_dict': self.discriminator.state_dict(),
155.             'optimizer_D_state_dict': self.optimizer_D.state_dict(),
156.         }, save_path)
157.         print(f"模型已保存至 {save_path}")
160. # ===========================
161. # 生成器
162. # ===========================
163. class Generator(nn.Module):
164.     def __init__(self, latent_dim=100, num_classes=4, img_channels=3):
165.         super(Generator, self).__init__()
166.         self.latent_dim = latent_dim
167.         self.label_emb = nn.Embedding(num_classes, num_classes)
168.         self.init_size = 8
169.         self.l1 = nn.Linear(latent_dim + num_classes, 256 * self.init_size * self.init_size)
171.         self.conv_blocks = nn.Sequential(
172.             nn.BatchNorm2d(256),
173.             nn.Upsample(scale_factor=2),
174.             nn.Conv2d(256, 128, 3, padding=1),
175.             nn.BatchNorm2d(128),
176.             nn.LeakyReLU(0.2, inplace=True),
178.             nn.Upsample(scale_factor=2),
179.             nn.Conv2d(128, 64, 3, padding=1),
180.             nn.BatchNorm2d(64),
181.             nn.LeakyReLU(0.2, inplace=True),
183.             nn.Upsample(scale_factor=2),
184.             nn.Conv2d(64, 32, 3, padding=1),
185.             nn.BatchNorm2d(32),
186.             nn.LeakyReLU(0.2, inplace=True),
188.             nn.Upsample(scale_factor=2),
189.             nn.Conv2d(32, img_channels, 3, padding=1),
190.             nn.Tanh()
191.         )
193.     def forward(self, noise, labels):
194.         labels = labels.to(self.label_emb.weight.device)
195.         label_embedding = self.label_emb(labels)
196.         x = torch.cat((noise, label_embedding), dim=1)
197.         x = self.l1(x).view(x.size(0), 256, self.init_size, self.init_size)
198.         return self.conv_blocks(x)
201. # ===========================
202. # 判别器
203. # ===========================
204. class Discriminator(nn.Module):
205.     def __init__(self, img_channels=3, num_classes=4):
206.         super(Discriminator, self).__init__()
207.         self.label_embedding = nn.Embedding(num_classes, img_channels)
209.         self.model = nn.Sequential(
210.             nn.Conv2d(img_channels * 2, 64, 4, stride=2, padding=1),
211.             nn.LeakyReLU(0.2, inplace=True),
212.             nn.Conv2d(64, 128, 4, stride=2, padding=1),
213.             nn.BatchNorm2d(128),
214.             nn.LeakyReLU(0.2, inplace=True),
216.             nn.Conv2d(128, 256, 4, stride=2, padding=1),
217.             nn.BatchNorm2d(256),
218.             nn.LeakyReLU(0.2, inplace=True),
220.             nn.Conv2d(256, 512, 4, stride=2, padding=1),
221.             nn.BatchNorm2d(512),
222.             nn.LeakyReLU(0.2, inplace=True)
223.         )
225.         self.output_layer = nn.Sequential(
226.             nn.Linear(512 * 8 * 8, 1),
227.             nn.Sigmoid()
228.         )
230.     def forward(self, img, labels):
231.         labels = labels.to(self.label_embedding.weight.device)
232.         label_embedding = self.label_embedding(labels).unsqueeze(2).unsqueeze(3)
233.         label_embedding = label_embedding.expand(-1, -1, img.size(2), img.size(3))
234.         x = torch.cat((img, label_embedding), dim=1)
235.         x = self.model(x).view(x.size(0), -1)
236.         return self.output_layer(x)

复制代码

2.3新建cGAN_trainer.py

1. import os
2. import torch
3. import argparse
4. from cGAN_net import cDCGAN
5. from utils import plot_loss, plot_result
6. import time
8. os.environ['OMP_NUM_THREADS'] = '1'
11. def main(args):
12.     # 初始化设备和训练参数
13.     device = torch.device(args.device if torch.cuda.is_available() else 'cpu')
14.     model = cDCGAN(data_root=args.data_root, batch_size=args.batch_size, device=device, latent_dim=args.latent_dim)
16.     # 添加学习率调度器
17.     scheduler_G = torch.optim.lr_scheduler.StepLR(model.optimizer_G, step_size=10, gamma=0.5)
18.     scheduler_D = torch.optim.lr_scheduler.StepLR(model.optimizer_D, step_size=10, gamma=0.5)
20.     start_epoch = 0
22.     # 如果有保存的模型，加载
23.     if args.load_model and os.path.exists(args.load_model):
24.         start_epoch = model.load_model(args.load_model) + 1
26.         # 恢复调度器状态
27.         scheduler_G_path = f"{args.load_model}_scheduler_G.pt"
28.         scheduler_D_path = f"{args.load_model}_scheduler_D.pt"
30.         if os.path.exists(scheduler_G_path) and os.path.exists(scheduler_D_path):
31.             scheduler_G.load_state_dict(torch.load(scheduler_G_path))
32.             scheduler_D.load_state_dict(torch.load(scheduler_D_path))
33.             print(f"成功恢复调度器状态：{scheduler_G_path}, {scheduler_D_path}")
34.         else:
35.             print("未找到调度器状态文件，使用默认调度器设置")
37.         print(f"从第 {start_epoch} 轮继续训练...")
39.     print(f"开始训练，从第 {start_epoch + 1} 轮开始...")
41.     # 创建保存路径
42.     os.makedirs(args.save_dir, exist_ok=True)
43.     os.makedirs(os.path.join(args.save_dir, 'log'), exist_ok=True)
45.     # 训练循环
46.     D_avg_losses, G_avg_losses = [], []
48.     for epoch in range(start_epoch, args.epochs):
49.         G_losses, D_losses, step = model.train_step(epoch, step=0, num_epochs=args.epochs)
52.         # 计算平均损失
53.         D_avg_loss = sum(D_losses) / len(D_losses) if D_losses else 0.0
54.         G_avg_loss = sum(G_losses) / len(G_losses) if G_losses else 0.0
55.         D_avg_losses.append(D_avg_loss)
56.         G_avg_losses.append(G_avg_loss)
58.         # 保存损失曲线图
59.         plot_loss(start_epoch, args.epochs, D_avg_losses, G_avg_losses, epoch + 1, save=True,
60.                   save_dir=os.path.join(args.save_dir, "log"))
62.         # 生成并保存图片
63.         labels = torch.tensor([0, 1, 2, 3]).to(device)
64.         if (epoch + 1) % args.save_freq == 0:  # 每隔一定轮次保存生成结果
65.             z = torch.randn(len(labels), args.latent_dim, device=device)  # 随机生成噪声
66.             plot_result(model.generator, z, labels, epoch + 1, save_dir=os.path.join(args.save_dir, 'log'))
68.         # 每10个epoch保存模型
69.         if (epoch + 1) % args.save_interval == 0:
70.             timestamp = int(time.time())
71.             save_path = os.path.join(args.save_dir, f"cgan_epoch_{epoch + 1}_{timestamp}.pth")
72.             model.save_model(epoch + 1, save_path)
73.             print(f"第 {epoch + 1} 轮的模型已保存，保存路径为 {save_path}")
75.         # 更新学习率调度器
76.         scheduler_G.step()
77.         scheduler_D.step()
79. if __name__ == "__main__":
80.     parser = argparse.ArgumentParser()
81.     parser.add_argument('--data_root', type=str, default='data/crop128', help="数据集根目录")
82.     parser.add_argument('--save_dir', type=str, default='./chkpt/cgan_model', help="保存模型的目录")
83.     parser.add_argument('--load_model', type=str, default=None, help="要加载的模型路径（可选）")
84.     parser.add_argument('--epochs', type=int, default=1000, help="训练的轮数")
85.     parser.add_argument('--save_interval', type=int, default=10, help="保存模型检查点的间隔（按轮数）")
86.     parser.add_argument('--batch_size', type=int, default=64, help="训练的批次大小")
87.     parser.add_argument('--device', type=str, default='cuda', help="使用的设备（如 cuda 或 cpu）")
88.     parser.add_argument('--latent_dim', type=int, default=100, help="生成器的潜在空间维度")
89.     parser.add_argument('--save_freq', type=int, default=1, help="每隔多少轮保存一次生成结果（默认: 1）")
90.     args = parser.parse_args()
91.     main(args)

复制代码

效果分析：

2.4中间效果可视化处理

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

1. import matplotlib.pyplot as plt
2. import numpy as np
3. import os
4. import torch
7. def denorm(x):
8.     out = (x + 1) / 2
9.     return out.clamp(0, 1)
12. def plot_loss(start_epoch, num_epochs, d_losses, g_losses, num_epoch, save=False, save_dir='celebA_cDCGAN_results/', show=False):
13.     """
14.     绘制损失函数曲线，从 start_epoch 到 num_epochs。
16.     Args:
17.         start_epoch: 起始轮次
18.         num_epochs: 总轮次
19.         d_losses: 判别器损失列表
20.         g_losses: 生成器损失列表
21.         num_epoch: 当前训练轮次
22.         save: 是否保存绘图
23.         save_dir: 保存路径
24.         show: 是否显示绘图
25.     """
26.     fig, ax = plt.subplots()
27.     ax.set_xlim(start_epoch, num_epochs)
28.     ax.set_ylim(0, max(np.max(g_losses), np.max(d_losses)) * 1.1)
30.     plt.xlabel(f'Epoch {num_epoch + 1}')
31.     plt.ylabel('Loss values')
33.     plt.plot(d_losses, label='Discriminator')
34.     plt.plot(g_losses, label='Generator')
35.     plt.legend()
37.     if save:
38.         if not os.path.exists(save_dir):
39.             os.makedirs(save_dir)
40.         save_fn = os.path.join(save_dir, f'cDCGAN_losses_epoch.png')
41.         plt.savefig(save_fn)
43.     if show:
44.         plt.show()
45.     else:
46.         plt.close()
49. def plot_result(generator, z, labels, epoch, save_dir=None, show=False):
50.     """
51.     生成并保存或显示生成的图片结果。
53.     Args:
54.         generator: 生成器模型
55.         z: 随机噪声张量
56.         labels: 标签张量
57.         epoch: 当前训练轮数
58.         save_dir: 保存图片的路径（可选）
59.         show: 是否显示生成的图片（可选）
60.     """
61.     # 调用生成器，生成图像
62.     generator.eval()  # 设置为评估模式
63.     with torch.no_grad():
64.         gen_images = generator(z, labels)  # 同时传入 z 和 labels
65.     generator.train()  # 恢复训练模式
67.     # 图像反归一化
68.     gen_images = denorm(gen_images)
70.     # 绘制图片
71.     fig, ax = plt.subplots(1, len(gen_images), figsize=(15, 15))
72.     for i in range(len(gen_images)):
73.         ax[i].imshow(gen_images[i].permute(1, 2, 0).cpu().numpy())  # 转换为可显示格式
74.         ax[i].axis('off')
76.     # 保存或显示图片
77.     if save_dir:
78.         os.makedirs(save_dir, exist_ok=True)
79.         save_path = os.path.join(save_dir, f'epoch_{epoch}.png')
80.         plt.savefig(save_path)
81.     if show:
82.         plt.show()
83.     plt.close(fig)

复制代码

执行 cGAN_trainer.py 文件，完成模子练习。

免责声明：如果侵犯了您的权益，请联系站长，我们会及时删除侵权内容，谢谢合作！更多信息从访问主页：qidao123.com:ToB企服之家，中国第一个企服评测及商务社交产业平台。