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()