当前位置: 首页 > ai >正文

GANs生成对抗网络生成手写数字的Pytorch实现

目录

一、第三方库导入

二、数据集准备

三、使用转置卷积的生成器

四、使用卷积的判别器

五、生成器生成图像

六、主程序

七、运行结果

7.1 生成器和判别器的损失函数图像

7.2 训练过程中生成器生成的图像

八、完整的pytorch代码


由于之前写gans的代码时,我的生成器和判别器不是使用的全连接网络就是卷积,但是无论这两种方法怎么组合,最后生成器生成的图像效果都很不好。因此最后我选择了生成器使用转置卷积,而判别器使用卷积,最后得到的生成图像确实效果比之前好很多了。

一、第三方库导入

import torch
import torch.nn as nn
import numpy as np
import matplotlib.pyplot as plt
plt.rcParams['font.sans-serif'] = ['SimHei']  # 设置中文字体
plt.rcParams['axes.unicode_minus'] = False  # 正常显示负号
from torchvision import transforms
import os
from PIL import Image
from torch.utils.data import Dataset, DataLoader

二、数据集准备

# 手写数字数据集
class MINISTDataset(Dataset):def __init__(self, files, root_dir, transform=None):self.files = filesself.root_dir = root_dirself.transform = transformself.labels = []for f in files:parts = f.split("_")p = parts[2].split(".")[0]self.labels.append(int(p))def __len__(self):return len(self.files)def __getitem__(self, idx):img_path = os.path.join(self.root_dir, self.files[idx])img = Image.open(img_path).convert("L")if self.transform:img = self.transform(img)label = self.labels[idx]return img, label

三、使用转置卷积的生成器

class Generator(nn.Module):def __init__(self, latent_dim=100):super().__init__()self.main = nn.Sequential(# 输入: latent_dim维噪声 -> 输出: 7x7x256nn.ConvTranspose2d(latent_dim, 256, kernel_size=7, stride=1, padding=0, bias=False),nn.BatchNorm2d(256),nn.ReLU(True),# 上采样: 7x7 -> 14x14nn.ConvTranspose2d(256, 128, kernel_size=4, stride=2, padding=1, bias=False),nn.BatchNorm2d(128),nn.ReLU(True),# 上采样: 14x14 -> 28x28nn.ConvTranspose2d(128, 64, kernel_size=4, stride=2, padding=1, bias=False),nn.BatchNorm2d(64),nn.ReLU(True),# 输出层: 28x28x1nn.ConvTranspose2d(64, 1, kernel_size=3, stride=1, padding=1, bias=False),nn.Tanh())def forward(self, x):# 将噪声重塑为 (batch_size, latent_dim, 1, 1)x = x.view(x.size(0), -1, 1, 1)return self.main(x)

四、使用卷积的判别器

class Discriminator(nn.Module):def __init__(self):super().__init__()self.main = nn.Sequential(# 输入: 1x28x28nn.Conv2d(1, 32, kernel_size=4, stride=2, padding=1),  # 输出: 32x14x14nn.LeakyReLU(0.2, inplace=True),nn.Dropout2d(0.3),nn.Conv2d(32, 64, kernel_size=4, stride=2, padding=1),  # 输出: 64x7x7nn.BatchNorm2d(64),nn.LeakyReLU(0.2, inplace=True),nn.Dropout2d(0.3),nn.Conv2d(64, 128, kernel_size=3, stride=1, padding=1),  # 输出: 128x7x7nn.BatchNorm2d(128),nn.LeakyReLU(0.2, inplace=True),nn.Dropout2d(0.3),nn.Flatten(),nn.Linear(128 * 7 * 7, 1),nn.Sigmoid())def forward(self, x):return self.main(x)

五、生成器生成图像

# 展示生成器生成的图像
def gen_img_plot(test_input, save_path):gen_imgs = gen(test_input).detach().cpu()gen_imgs = gen_imgs.view(-1, 28, 28)plt.figure(figsize=(4, 4))for i in range(16):plt.subplot(4, 4, i + 1)plt.imshow(gen_imgs[i], cmap="gray")plt.axis("off")plt.savefig(save_path, dpi=300)plt.close()

六、主程序

if __name__ == "__main__":# 对数据做归一化处理transforms = transforms.Compose([transforms.Resize((28, 28)),transforms.ToTensor(),transforms.Normalize((0.5,), (0.5,))])# 路径base_dir = 'C:\\Users\\Administrator\\PycharmProjects\\CNN'train_dir = os.path.join(base_dir, "minist_train")# 获取文件夹里图像的名称train_files = [f for f in os.listdir(train_dir) if f.endswith('.jpg')]# 创建数据集和数据加载器train_dataset = MINISTDataset(train_files, train_dir, transform=transforms)train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)# 参数epochs = 50lr = 0.0002# 初始化模型的优化器和损失函数gen = Generator()dis = Discriminator()d_optim = torch.optim.Adam(dis.parameters(), lr=lr, betas=(0.5, 0.999))  # 判别器的优化器g_optim = torch.optim.Adam(gen.parameters(), lr=lr, betas=(0.5, 0.999))  # 生成器的优化器loss_fn = torch.nn.BCELoss()  # 二分类交叉熵损失函数# 记录lossD_loss = []G_loss = []# 训练for epoch in range(epochs):d_epoch_loss = 0g_epoch_loss = 0count = len(train_loader)  # 返回批次数for step, (img, _) in enumerate(train_loader):# 每个批次的大小size = img.size(0)random_noise = torch.randn(size, 100)# 判别器训练d_optim.zero_grad()real_output = dis(img)d_real_loss = loss_fn(real_output, torch.ones_like(real_output))# d_real_loss.backward()gen_img = gen(random_noise)gen_img = gen_img.view(size, 1, 28, 28)fake_output = dis(gen_img.detach())d_fake_loss = loss_fn(fake_output, torch.zeros_like(fake_output))# d_fake_loss.backward()d_loss = (d_real_loss + d_fake_loss) / 2d_loss.backward()d_optim.step()# 生成器的训练g_optim.zero_grad()fake_output = dis(gen_img)g_loss = loss_fn(fake_output, torch.ones_like(fake_output))g_loss.backward()g_optim.step()# 计算在一个epoch里面所有的g_loss和d_losswith torch.no_grad():d_epoch_loss += d_lossg_epoch_loss += g_loss# 计算平均损失值with torch.no_grad():d_epoch_loss = d_epoch_loss / countg_epoch_loss = g_epoch_loss / countD_loss.append(d_epoch_loss.item())G_loss.append(g_epoch_loss.item())print("Epoch:", epoch, "  D loss:", d_epoch_loss.item(), "  G Loss:", g_epoch_loss.item())# 每隔2个epoch绘制生成器生成的图像if (epoch + 1) % 2 == 0:test_input = torch.randn(16, 100)name = f"gen_img_{epoch}.jpg"save_path = os.path.join('C:\\Users\\Administrator\\PycharmProjects\\CNN\\gen_img_11', name)gen_img_plot(test_input, save_path)# 绘制损失曲线图plt.figure(figsize=(12, 6))plt.plot(D_loss, label="判别器", color="tomato")plt.plot(G_loss, label="生成器", color="orange")plt.xlabel("epoch")plt.ylabel("loss")plt.title("生成器和判别器的损失曲线图")plt.legend()plt.grid()plt.savefig("C:\\Users\\Administrator\\PycharmProjects\\CNN\\gen_dis_loss_11.jpg", dpi=300, bbox_inches="tight")plt.close()

七、运行结果

7.1 生成器和判别器的损失函数图像

7.2 训练过程中生成器生成的图像

这里只展示一部分

gen_img_1.jpg

gen_img_25.jpg

gen_img_49.jpg

八、完整的pytorch代码

import torch
import torch.nn as nn
import numpy as np
import matplotlib.pyplot as plt
plt.rcParams['font.sans-serif'] = ['SimHei']  # 设置中文字体
plt.rcParams['axes.unicode_minus'] = False  # 正常显示负号
from torchvision import transforms
import os
from PIL import Image
from torch.utils.data import Dataset, DataLoader# 手写数字数据集
class MINISTDataset(Dataset):def __init__(self, files, root_dir, transform=None):self.files = filesself.root_dir = root_dirself.transform = transformself.labels = []for f in files:parts = f.split("_")p = parts[2].split(".")[0]self.labels.append(int(p))def __len__(self):return len(self.files)def __getitem__(self, idx):img_path = os.path.join(self.root_dir, self.files[idx])img = Image.open(img_path).convert("L")if self.transform:img = self.transform(img)label = self.labels[idx]return img, label# 改进的生成器(使用转置卷积)
class Generator(nn.Module):def __init__(self, latent_dim=100):super().__init__()self.main = nn.Sequential(# 输入: latent_dim维噪声 -> 输出: 7x7x256nn.ConvTranspose2d(latent_dim, 256, kernel_size=7, stride=1, padding=0, bias=False),nn.BatchNorm2d(256),nn.ReLU(True),# 上采样: 7x7 -> 14x14nn.ConvTranspose2d(256, 128, kernel_size=4, stride=2, padding=1, bias=False),nn.BatchNorm2d(128),nn.ReLU(True),# 上采样: 14x14 -> 28x28nn.ConvTranspose2d(128, 64, kernel_size=4, stride=2, padding=1, bias=False),nn.BatchNorm2d(64),nn.ReLU(True),# 输出层: 28x28x1nn.ConvTranspose2d(64, 1, kernel_size=3, stride=1, padding=1, bias=False),nn.Tanh())def forward(self, x):# 将噪声重塑为 (batch_size, latent_dim, 1, 1)x = x.view(x.size(0), -1, 1, 1)return self.main(x)# 改进的判别器(使用深度卷积网络)
class Discriminator(nn.Module):def __init__(self):super().__init__()self.main = nn.Sequential(# 输入: 1x28x28nn.Conv2d(1, 32, kernel_size=4, stride=2, padding=1),  # 输出: 32x14x14nn.LeakyReLU(0.2, inplace=True),nn.Dropout2d(0.3),nn.Conv2d(32, 64, kernel_size=4, stride=2, padding=1),  # 输出: 64x7x7nn.BatchNorm2d(64),nn.LeakyReLU(0.2, inplace=True),nn.Dropout2d(0.3),nn.Conv2d(64, 128, kernel_size=3, stride=1, padding=1),  # 输出: 128x7x7nn.BatchNorm2d(128),nn.LeakyReLU(0.2, inplace=True),nn.Dropout2d(0.3),nn.Flatten(),nn.Linear(128 * 7 * 7, 1),nn.Sigmoid())def forward(self, x):return self.main(x)# 展示生成器生成的图像
def gen_img_plot(test_input, save_path):gen_imgs = gen(test_input).detach().cpu()gen_imgs = gen_imgs.view(-1, 28, 28)plt.figure(figsize=(4, 4))for i in range(16):plt.subplot(4, 4, i + 1)plt.imshow(gen_imgs[i], cmap="gray")plt.axis("off")plt.savefig(save_path, dpi=300)plt.close()if __name__ == "__main__":# 对数据做归一化处理transforms = transforms.Compose([transforms.Resize((28, 28)),transforms.ToTensor(),transforms.Normalize((0.5,), (0.5,))])# 路径base_dir = 'C:\\Users\\Administrator\\PycharmProjects\\CNN'train_dir = os.path.join(base_dir, "minist_train")# 获取文件夹里图像的名称train_files = [f for f in os.listdir(train_dir) if f.endswith('.jpg')]# 创建数据集和数据加载器train_dataset = MINISTDataset(train_files, train_dir, transform=transforms)train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)# 参数epochs = 50lr = 0.0002# 初始化模型的优化器和损失函数gen = Generator()dis = Discriminator()d_optim = torch.optim.Adam(dis.parameters(), lr=lr, betas=(0.5, 0.999))  # 判别器的优化器g_optim = torch.optim.Adam(gen.parameters(), lr=lr, betas=(0.5, 0.999))  # 生成器的优化器loss_fn = torch.nn.BCELoss()  # 二分类交叉熵损失函数# 记录lossD_loss = []G_loss = []# 训练for epoch in range(epochs):d_epoch_loss = 0g_epoch_loss = 0count = len(train_loader)  # 返回批次数for step, (img, _) in enumerate(train_loader):# 每个批次的大小size = img.size(0)random_noise = torch.randn(size, 100)# 判别器训练d_optim.zero_grad()real_output = dis(img)d_real_loss = loss_fn(real_output, torch.ones_like(real_output))# d_real_loss.backward()gen_img = gen(random_noise)gen_img = gen_img.view(size, 1, 28, 28)fake_output = dis(gen_img.detach())d_fake_loss = loss_fn(fake_output, torch.zeros_like(fake_output))# d_fake_loss.backward()d_loss = (d_real_loss + d_fake_loss) / 2d_loss.backward()d_optim.step()# 生成器的训练g_optim.zero_grad()fake_output = dis(gen_img)g_loss = loss_fn(fake_output, torch.ones_like(fake_output))g_loss.backward()g_optim.step()# 计算在一个epoch里面所有的g_loss和d_losswith torch.no_grad():d_epoch_loss += d_lossg_epoch_loss += g_loss# 计算平均损失值with torch.no_grad():d_epoch_loss = d_epoch_loss / countg_epoch_loss = g_epoch_loss / countD_loss.append(d_epoch_loss.item())G_loss.append(g_epoch_loss.item())print("Epoch:", epoch, "  D loss:", d_epoch_loss.item(), "  G Loss:", g_epoch_loss.item())# 每隔2个epoch绘制生成器生成的图像if (epoch + 1) % 2 == 0:test_input = torch.randn(16, 100)name = f"gen_img_{epoch}.jpg"save_path = os.path.join('C:\\Users\\Administrator\\PycharmProjects\\CNN\\gen_img_11', name)gen_img_plot(test_input, save_path)# 绘制损失曲线图plt.figure(figsize=(12, 6))plt.plot(D_loss, label="判别器", color="tomato")plt.plot(G_loss, label="生成器", color="orange")plt.xlabel("epoch")plt.ylabel("loss")plt.title("生成器和判别器的损失曲线图")plt.legend()plt.grid()plt.savefig("C:\\Users\\Administrator\\PycharmProjects\\CNN\\gen_dis_loss_11.jpg", dpi=300, bbox_inches="tight")plt.close()
http://www.xdnf.cn/news/17986.html

相关文章:

  • 原型和原型链的问题
  • mac电脑开发嵌入式基于Clion(stm32CubeMX)
  • ThinkPHP8学习篇(三):控制器
  • 《解构WebSocket断网重连:指数退避算法的前端工业级实践指南》
  • pair之于vector、queue(vector<pair<int,int>>)
  • Yolov模型的演变
  • K8S集群环境搭建
  • 【LeetCode 热题 100】(八)二叉树
  • 数据结构——栈和队列oj练习
  • 深度解析 Spring Bean 生命周期
  • 【网络安全】Webshell的绕过——绕过动态检测引擎WAF-缓存绕过(Hash碰撞)
  • 《P4180 [BJWC2010] 严格次小生成树》
  • MySQL 插入数据提示字段超出范围?一招解决 DECIMAL 类型踩坑
  • 安卓11 12系统修改定制化_____修改运营商版本安装特定应用时的默认规则
  • 机器学习相关算法:回溯算法 贪心算法 回归算法(线性回归) 算法超参数 多项式时间 朴素贝叶斯分类算法
  • 一文速通Python并行计算:14 Python异步编程-协程的管理和调度
  • C语言:文件操作详解
  • 后量子密码算法SLH-DSA介绍及开源代码实现
  • Java8~Java21重要新特性
  • C++ 最短路Dijkstra
  • CodeBuddy IDE深度体验:AI驱动的全栈开发新时代
  • Maven下载和配置-IDEA使用
  • 【算法】——力扣hot100常用算法技巧
  • 使用IntersectionObserver实现页面右侧运营位区域固定,和页面列表数据分页加载
  • JetPack系列教程(七):Palette——让你的APP色彩“飞”起来!
  • 【大语言模型 02】多头注意力深度剖析:为什么需要多个头
  • 后量子密码算法ML-DSA介绍及开源代码实现
  • 【DL学习笔记】常用数据集总结
  • 微服务架构实战指南:从单体应用到云原生的蜕变之路
  • 56. 合并区间