pytorch搭建GAN网络生成mnist假图片
效果如下图所示:
完整代码
import os
import torch
import torchvision
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
from torchvision import datasets
from torchvision.utils import save_image
batch_size = 100
epoch_num = 30
lr = 0.0002
input_dim = 100
class Generator(nn.Module):
def __init__(self, input_dim):
super(Generator, self).__init__()
self.fc1 = nn.Linear(input_dim, 56 * 56)
self.br = nn.Sequential(
nn.BatchNorm2d(1),
nn.ReLU(True)
)
self.conv1 = nn.Sequential(
nn.Conv2d(1, 50, 3, stride=1, padding=1),
nn.BatchNorm2d(50),
nn.ReLU(True)
)
self.conv2 = nn.Sequential(
nn.Conv2d(50, 25, 3, stride=1, padding=1),
nn.BatchNorm2d(25),
nn.ReLU(True)
)
self.conv3 = nn.Sequential(
nn.Conv2d(25, 1, 2, stride=2),
nn.Tanh()
)
def forward(self, x):
x = self.fc1(x)
x = x.view(-1, 1, 56, 56)
x = self.br(x)
x = self.conv1(x)
x = self.conv2(x)
output = self.conv3(x)
return output
class Discriminator(nn.Module):
def __init__(self):
super(Discriminator, self).__init__()
self.conv1 = nn.Sequential(
nn.Conv2d(1, 32, 5, stride=1, padding=2),
nn.LeakyReLU(0.2,True)
)
self.pl1 = nn.AvgPool2d(2, stride=2)
self.conv2 = nn.Sequential(
nn.Conv2d(32, 64, 5, stride=1, padding=2),
nn.LeakyReLU(0.2,True)
)
self.pl2 = nn.AvgPool2d(2, stride=2)
self.fc1 = nn.Sequential(
nn.Linear(64 * 7 * 7, 1024),
nn.LeakyReLU(0.2,True)
)
self.fc2 = nn.Sequential(
nn.Linear(1024, 1),
nn.Sigmoid()
)
def forward(self, x):
x = self.conv1(x)
x = self.pl1(x)
x = self.conv2(x)
x = self.pl2(x)
x = x.view(x.shape[0], -1)
x = self.fc1(x)
output = self.fc2(x)
return output
def G_train(input_dim):
G_optimizer.zero_grad()
noise = torch.randn(batch_size, input_dim).to(device)
real_label = torch.ones(batch_size).to(device)
fake_img = G(noise)
D_output = D(fake_img)
G_loss = criterion(D_output, real_label)
G_loss.backward()
G_optimizer.step()
return G_loss.data.item()
def D_train(real_img, input_dim):
D_optimizer.zero_grad()
real_label = torch.ones(real_img.shape[0]).to(device)
D_output = D(real_img)
D_real_loss = criterion(D_output, real_label)
noise = torch.randn(batch_size, input_dim, requires_grad=False).to(device)
fake_label = torch.zeros(batch_size).to(device)
fake_img = G(noise)
D_output = D(fake_img)
D_fake_loss = criterion(D_output, fake_label)
D_loss = D_real_loss + D_fake_loss
D_loss.backward()
D_optimizer.step()
return D_loss.data.item()
def save_img(img, img_name):
img = 0.5 * (img + 1)
img = img.clamp(0, 1)
save_image(img, "./imgs/" + img_name)
if __name__ == "__main__":
if not os.path.exists("./checkpoint"):
os.makedirs("./checkpoint")
if not os.path.exists("./imgs"):
os.makedirs("./imgs")
if not os.path.exists("./mnist"):
os.makedirs("./mnist")
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
# 加载数据
train_dataset = datasets.MNIST(root='./mnist/', train=True, transform=torchvision.transforms.ToTensor(),
download=True)
train_loader = torch.utils.data.DataLoader(dataset=train_dataset, batch_size=batch_size, shuffle=True)
# 构建生成器和判别器网络
if os.path.exists('./checkpoint/Generator.pkl') and os.path.exists('./checkpoint/Discriminator.pkl'):
G=torch.load("./checkpoint/Generator.pkl").to(device)
D=torch.load("./checkpoint/Discriminator.pkl").to(device)
else:
G = Generator(input_dim).to(device)
D = Discriminator().to(device)
# 指明损失函数和优化器
criterion = nn.BCELoss()
G_optimizer = optim.Adam(G.parameters(), lr=lr)
D_optimizer = optim.Adam(D.parameters(), lr=lr)
print("Training...........")
for epoch in range(1, epoch_num + 1):
print("epoch: ", epoch)
for batch, (x, _) in enumerate(train_loader):
# 对判别器和生成器分别进行训练,注意顺序不能反
D_loss=D_train(x.to(device), input_dim)
G_loss=G_train(input_dim)
#if batch % 20 == 0:
print("[ %d / %d ] g_loss: %.6f d_loss: %.6f" % (batch, 600, float(G_loss), float(D_loss)))
if batch % 50 == 0:
fake_img = torch.randn(128, input_dim).to(device)
fake_img = G(fake_img)
save_img(fake_img, "img_" + str(epoch) + "_" + str(batch) + ".png")
# 保存模型
torch.save(G, "./checkpoint/Generator.pkl")
torch.save(D, "./checkpoint/Discriminator.pkl")