【专知-PyTorch手把手深度学习教程04】GAN快速理解与PyTorch实现: 图文+代码

import torch
import torch.nn as nn
from torchvision import datasets
from torchvision import transforms
from torchvision.utils import save_image
from torch.autograd import Variable

def get_variable(x):
    x = Variable(x)
    return x.cuda() if torch.cuda.is_available() else x

def denorm(x):
    out = (x + 1) / 2
    return out.clamp(0, 1)

transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize(mean=(0.5, 0.5, 0.5),
                         std=(0.5, 0.5, 0.5))])
mnist = datasets.MNIST(root='./mnist/',
                       train=True,
                       transform=transform,
                       download=True)
data_loader = torch.utils.data.DataLoader(dataset=mnist,
                                          batch_size=100,
                                          shuffle=True)
# 判别器
D = nn.Sequential(
    nn.Linear(784, 256),
    nn.LeakyReLU(0.2),
    nn.Linear(256, 256),
    nn.LeakyReLU(0.2),
    nn.Linear(256, 1),
    nn.Sigmoid())

# 生成器
G = nn.Sequential(
    nn.Linear(64, 256),
    nn.LeakyReLU(0.2),
    nn.Linear(256, 256),
    nn.LeakyReLU(0.2),
    nn.Linear(256, 784),
    nn.Tanh())

if torch.cuda.is_available():
    D.cuda()
    G.cuda()

loss_func = nn.BCELoss()
d_optimizer = torch.optim.Adam(D.parameters(), lr=0.0003)
g_optimizer = torch.optim.Adam(G.parameters(), lr=0.0003)

for epoch in range(200):
    for i, (images, _) in enumerate(data_loader):

        batch_size = images.size(0)
        # reshape 成 (batch_size, 28*28)
        images = get_variable(images.view(batch_size, -1))
        real_labels = get_variable(torch.ones(batch_size))  # 真实数据 label 为1
        fake_labels = get_variable(torch.zeros(batch_size))  # 假数据 label 为0

        # ============= Train the discriminator =============#
        # 判别真实数据,计算损失
        outputs = D(images)
        d_loss_real = loss_func(outputs, real_labels)
        real_score = outputs

        # 生成假数据
        z = get_variable(torch.randn(batch_size, 64))
        fake_images = G(z)

        # 判别生成的数据,计算损失
        outputs = D(fake_images)
        d_loss_fake = loss_func(outputs, fake_labels)
        fake_score = outputs

        # 优化判别器
        d_loss = d_loss_real + d_loss_fake
        D.zero_grad()
        d_loss.backward()
        d_optimizer.step()

        # =============== Train the generator ===============#
        # 生成假数据
        z = get_variable(torch.randn(batch_size, 64))
        fake_images = G(z)
        # 用判别器计算损失
        outputs = D(fake_images)
        g_loss = loss_func(outputs, real_labels)

        # 优化生成器
        D.zero_grad()
        G.zero_grad()
        g_loss.backward()
        g_optimizer.step()

        if (i + 1) % 300 == 0:
            print('Epoch [%d/%d], Step[%d/%d], d_loss: %.4f, '
                  'g_loss: %.4f, 真实数据平均得分: %.2f, 假数据平均得分: %.2f'
                  % (epoch, 200, i + 1, 600, d_loss.data[0], g_loss.data[0],
                     real_score.data.mean(), fake_score.data.mean()))

    # 保存一下真实数据
    if (epoch + 1) == 1:
        images = images.view(images.size(0), 1, 28, 28)
        save_image(denorm(images.data), './mnist/real_images.png')

    # 保存生成数据
    fake_images = fake_images.view(fake_images.size(0), 1, 28, 28)
    save_image(denorm(fake_images.data), './mnist/fake_images-%d.png' % (epoch + 1))

# 保存模型参数
torch.save(G.state_dict(), './generator.pkl')
torch.save(D.state_dict(), './discriminator.pkl')