GAN生成对抗网络：花卉生成

简介

本篇文章利用pytorch搭建GAN生成对抗网络实现花卉生成的任务

一、GAN生成对抗网络基础知识

关于GAN生成对抗网络的基础知识以下文章有详细讲解，可供参考：
GAN（生成对抗网络）的系统全面介绍（醍醐灌顶）

二、数据集介绍

本文使用花卉数据集，该数据集包含了4317张图片，包含雏菊、蒲公英、玫瑰、向日葵、郁金香五种花卉，我已将数据集拆分为训练集和测试集两部分，本文仅使用了训练集部分，以下是数据集目录：

数据集已放于以下链接，有需要可自行下载
花卉数据集

三、代码实现

参数设置

step1.参数continue_train：是否继续训练
step2.参数dir：训练集路径
step3.参数batch_size：单次训练图片量
step4.参数device：使用GPU
step5.参数epochs：训练周期
step6.参数generator_num：每k轮训练一次生成器
step7.参数discriminator_num：每k轮训练一次判别器

if __name__ == '__main__':parser = argparse.ArgumentParser()parser.add_argument('--continue_train', type=bool, default=False, help='continue training')parser.add_argument('--dir', type=str, default='./flowers/train', help='dataset path')parser.add_argument('--batch_size', type=int, default=50, help='batch size')parser.add_argument('--device', type=int, default=0, help='GPU id')parser.add_argument('--epochs', type=int, default=200, help='train epochs')parser.add_argument('--generator_num', type=int, default=5, help='train generator every k epochs')parser.add_argument('--discriminator_num', type=int, default=1, help='train discriminator every k epochs')args = parser.parse_args()main(args)

数据处理

step1.定义训练集中图像输入判别器前的transform操作
step2.准备Dataset与Dataloader

    transform = transforms.Compose([transforms.Resize((96, 96)),  # 将图片resize至 96 * 96transforms.ToTensor(),  # 转换为张量transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])data_set = datasets.ImageFolder(root=args.dir, transform=transform)data_loader = dataloader.DataLoader(dataset=data_set, batch_size=args.batch_size, num_workers=4, shuffle=True, drop_last=True)print('already load data...')

搭建网络

step1.生成器使用反卷积，最终输出3 * 96 * 96大小的图片，且像素值$\in [-1,1]$
step2.生成器使用卷积，最终输出判别为真的概率

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),       # 512 × 4 × 4        (1-1)*1+1*(4-1)+0+1 = 4nn.ConvTranspose2d(512, 256, kernel_size=4, stride=2, padding=1, bias=False),nn.BatchNorm2d(256),nn.ReLU(True),      # 256 × 8 × 8     (4-1)*2-2*1+1*(4-1)+0+1 = 8nn.ConvTranspose2d(256, 128, kernel_size=4, stride=2, padding=1, bias=False),nn.BatchNorm2d(128),nn.ReLU(True),  # 128 × 16 × 16nn.ConvTranspose2d(128, 64, kernel_size=4, stride=2, padding=1, bias=False),nn.BatchNorm2d(64),nn.ReLU(True),      # 64 × 32 × 32nn.ConvTranspose2d(64, 3, kernel_size=5, stride=3, padding=1, bias=False),nn.Tanh()       # 3 * 96 * 96)def forward(self, input):return self.main(input)class Discriminator(nn.Module):def __init__(self):super(Discriminator,self).__init__()self.main = nn.Sequential(nn.Conv2d(3, 64, kernel_size=5, stride=3, padding=1, bias=False),nn.LeakyReLU(0.2, inplace=True),        # 64 * 32 * 32nn.Conv2d(64, 128, kernel_size=4, stride=2, padding=1, bias=False),nn.BatchNorm2d(128),nn.LeakyReLU(0.2, inplace=True),         # 128 * 16 * 16nn.Conv2d(128, 256, kernel_size=4, stride=2, padding=1, bias=False),nn.BatchNorm2d(256),nn.LeakyReLU(0.2, inplace=True),  # 256 * 8 * 8nn.Conv2d(256, 512, kernel_size=4, stride=2, padding=1, bias=False),nn.BatchNorm2d(512),nn.LeakyReLU(0.2, inplace=True),  # 512 * 4 * 4nn.Conv2d(512, 1, kernel_size=4, stride=1, padding=0, bias=False),nn.Sigmoid()        # 输出一个概率)def forward(self, input):return self.main(input).view(-1)

定义优化器与损失函数

step1.生成器与判别器的优化器都使用Adam
step2.将损失函数使用二元交叉熵损失

    optimizer_G = torch.optim.Adam(model_G.parameters(), lr=2e-4, betas=(0.5, 0.999))optimizer_D = torch.optim.Adam(model_D.parameters(), lr=2e-4, betas=(0.5, 0.999))loss = nn.BCELoss()print('already prepared optimizer and loss_function...')

训练网络

每discriminator_num轮：
step1.输入真图片让判别器鉴别
step2.生成器利用随机噪声生成图片，并让判别器鉴别
step3.计算判别器损失（真鉴别为真，假鉴别为假），反向传播后更新判别器参数
每generator_num轮：
step4.生成器利用随机噪声生成图片，并让判别器鉴别
step5.计算生成器损失（假鉴别为真），反向传播后更新生成器参数
step6.每100轮保存一次结果

    print('start training...')for epoch in range(args.epochs):print('epoch:{}'.format(epoch + 1))for i, data in enumerate(data_loader):if (i + 1) % args.discriminator_num == 0:optimizer_D.zero_grad()real_img = data[0]batchsize = len(real_img)real_img = real_img.cuda(args.device)out_D_real = model_D(real_img)real_labels = torch.ones(batchsize).cuda(args.device)loss_D_real = loss(out_D_real, real_labels)loss_D_real.backward()noise = torch.randn(args.batch_size, 100, 1, 1).cuda(args.device)fake_img = model_G(noise)out_D_fake = model_D(fake_img)fake_labels = torch.zeros(batchsize).cuda(args.device)loss_D_fake = loss(out_D_fake, fake_labels)loss_D_fake.backward()optimizer_D.step()if (i + 1) % args.generator_num == 0:optimizer_G.zero_grad()real_img = data[0]batchsize = len(real_img)noise = torch.randn(args.batch_size, 100, 1, 1).cuda(args.device)fake_img = model_G(noise)out_D_fake = model_D(fake_img)real_labels = torch.ones(batchsize).cuda(args.device)loss_G = loss(out_D_fake, real_labels)loss_G.backward()optimizer_G.step()if (epoch + 1) % 100 == 0:fix_noise = torch.randn(40, 100, 1, 1).cuda(args.device)final_img = model_G(fix_noise)final_img = final_img * 0.5 + 0.5final_img = final_img.cpu()plt.figure(1)for i in range(40):img = final_img[i].detach().numpy()plt.subplot(5, 8, i+1)plt.imshow(np.transpose(img, (1, 2, 0)))plt.savefig("./outcome/{}.png".format(epoch + 1))plt.show()print('end training...')

保存网络

    torch.save(model_G.state_dict(), './generator.pt')torch.save(model_D.state_dict(), './discriminator.pt')print('already saved model...')

结果展示

训练3000轮后得到结果如下：

总结

以上就是利用生成对抗网络实现图像生成的介绍，完整代码如下：

import argparse
import torchvision.datasets as datasets
import torch.utils.data.dataloader as dataloader
import torchvision.transforms as transforms
import torch.nn as nn
import torch
import numpy as np
import matplotlib.pyplot as pltclass 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),       # 512 × 4 × 4        (1-1)*1+1*(4-1)+0+1 = 4nn.ConvTranspose2d(512, 256, kernel_size=4, stride=2, padding=1, bias=False),nn.BatchNorm2d(256),nn.ReLU(True),      # 256 × 8 × 8     (4-1)*2-2*1+1*(4-1)+0+1 = 8nn.ConvTranspose2d(256, 128, kernel_size=4, stride=2, padding=1, bias=False),nn.BatchNorm2d(128),nn.ReLU(True),  # 128 × 16 × 16nn.ConvTranspose2d(128, 64, kernel_size=4, stride=2, padding=1, bias=False),nn.BatchNorm2d(64),nn.ReLU(True),      # 64 × 32 × 32nn.ConvTranspose2d(64, 3, kernel_size=5, stride=3, padding=1, bias=False),nn.Tanh()       # 3 * 96 * 96)def forward(self, input):return self.main(input)class Discriminator(nn.Module):def __init__(self):super(Discriminator,self).__init__()self.main = nn.Sequential(nn.Conv2d(3, 64, kernel_size=5, stride=3, padding=1, bias=False),nn.LeakyReLU(0.2, inplace=True),        # 64 * 32 * 32nn.Conv2d(64, 128, kernel_size=4, stride=2, padding=1, bias=False),nn.BatchNorm2d(128),nn.LeakyReLU(0.2, inplace=True),         # 128 * 16 * 16nn.Conv2d(128, 256, kernel_size=4, stride=2, padding=1, bias=False),nn.BatchNorm2d(256),nn.LeakyReLU(0.2, inplace=True),  # 256 * 8 * 8nn.Conv2d(256, 512, kernel_size=4, stride=2, padding=1, bias=False),nn.BatchNorm2d(512),nn.LeakyReLU(0.2, inplace=True),  # 512 * 4 * 4nn.Conv2d(512, 1, kernel_size=4, stride=1, padding=0, bias=False),nn.Sigmoid()        # 输出一个概率)def forward(self, input):return self.main(input).view(-1)def main(args):transform = transforms.Compose([transforms.Resize((96, 96)),  # 将图片resize至 96 * 96transforms.ToTensor(),  # 转换为张量transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])data_set = datasets.ImageFolder(root=args.dir, transform=transform)data_loader = dataloader.DataLoader(dataset=data_set, batch_size=args.batch_size, num_workers=4, shuffle=True, drop_last=True)print('already load data...')model_G = Generator()model_D = Discriminator()if args.continue_train == True:model_G.load_state_dict(torch.load('./generator.pt'))model_D.load_state_dict(torch.load('./discriminator.pt'))model_G.train()model_D.train()print('already prepared model...')optimizer_G = torch.optim.Adam(model_G.parameters(), lr=2e-4, betas=(0.5, 0.999))optimizer_D = torch.optim.Adam(model_D.parameters(), lr=2e-4, betas=(0.5, 0.999))loss = nn.BCELoss()print('already prepared optimizer and loss_function...')if torch.cuda.is_available() == True:model_G.cuda(args.device)model_D.cuda(args.device)loss.cuda(args.device)print('already in GPU...')print('start training...')for epoch in range(args.epochs):print('epoch:{}'.format(epoch + 1))for i, data in enumerate(data_loader):if (i + 1) % args.discriminator_num == 0:optimizer_D.zero_grad()real_img = data[0]batchsize = len(real_img)real_img = real_img.cuda(args.device)out_D_real = model_D(real_img)real_labels = torch.ones(batchsize).cuda(args.device)loss_D_real = loss(out_D_real, real_labels)loss_D_real.backward()noise = torch.randn(args.batch_size, 100, 1, 1).cuda(args.device)fake_img = model_G(noise)out_D_fake = model_D(fake_img)fake_labels = torch.zeros(batchsize).cuda(args.device)loss_D_fake = loss(out_D_fake, fake_labels)loss_D_fake.backward()optimizer_D.step()if (i + 1) % args.generator_num == 0:optimizer_G.zero_grad()real_img = data[0]batchsize = len(real_img)noise = torch.randn(args.batch_size, 100, 1, 1).cuda(args.device)fake_img = model_G(noise)out_D_fake = model_D(fake_img)real_labels = torch.ones(batchsize).cuda(args.device)loss_G = loss(out_D_fake, real_labels)loss_G.backward()optimizer_G.step()if (epoch + 1) % 10 == 0:fix_noise = torch.randn(40, 100, 1, 1).cuda(args.device)final_img = model_G(fix_noise)final_img = final_img * 0.5 + 0.5final_img = final_img.cpu()plt.figure(1)for i in range(40):img = final_img[i].detach().numpy()plt.subplot(5, 8, i+1)plt.imshow(np.transpose(img, (1, 2, 0)))plt.savefig("./outcome/{}.png".format(epoch + 1))plt.show()print('end training...')torch.save(model_G.state_dict(), './generator.pt')torch.save(model_D.state_dict(), './discriminator.pt')print('already saved model...')if __name__ == '__main__':parser = argparse.ArgumentParser()parser.add_argument('--continue_train', type=bool, default=False, help='continue training')parser.add_argument('--dir', type=str, default='./flowers/train', help='dataset path')parser.add_argument('--batch_size', type=int, default=50, help='batch size')parser.add_argument('--device', type=int, default=0, help='GPU id')parser.add_argument('--epochs', type=int, default=3000, help='train epochs')parser.add_argument('--generator_num', type=int, default=5, help='train generator every k epochs')parser.add_argument('--discriminator_num', type=int, default=1, help='train discriminator every k epochs')args = parser.parse_args()main(args)