第一个成功GAN项目: 生成手写数字(MNIST)

其实应该说是第1.5个项目,增加了独热向量影响判别器和生成器,使其能针对行输出目标数字

训练大概15个epoch,效果还可以,就是训练的时候生成器的学习率少写了个0(应该是0.0001

但是结果还不错,之后试着改了下学习率重新训练但是效果没有之前手误的好,生成器的loss一直比较高,所以改回来了

关于训练时间,我电脑纯cpu跑大概20-s/epoch,可以说是非常快奥(batch=64)

后来让朋友用gpu跑,我也用colab跑了一下,甚至比我cpu跑的还慢,不知道什么原因(可能是每次训练生成张量再发送到gpu占大部分?

代码(jupyter, pytorch, py3.9)

#%%
import random

import matplotlib.pyplot as plt
import numpy as np
import pandas
import torch
from torch import nn
from torch.utils.data import DataLoader
from torchvision import transforms, datasets
#%%
class D(nn.Module):
    def __init__(self):
        super(D, self).__init__()
        self.model = nn.Sequential(
            nn.Linear(784 + 10, 200),
            nn.LeakyReLU(0.02),
            nn.LayerNorm(200),
            nn.Linear(200, 1),
            nn.Sigmoid()
        )

        self.optimiser = torch.optim.Adam(self.parameters(), lr=0.0001)
        self.counter = 0
        self.progress = []
        self.loss_func = nn.BCELoss()

    def forward(self, inputs, label_tensors):
        return self.model(torch.cat((inputs.view(-1, 784), label_tensors), 1))

    def nn_train(self, inputs, label_tensors, targets):
        outputs = self(inputs, label_tensors)
        loss = self.loss_func(outputs, targets)
        self.counter += 1
        if not self.counter % 10:
            self.progress.append(loss.item())
        self.optimiser.zero_grad()
        loss.backward()
        self.optimiser.step()

    def plot_progress(self):
        df = pandas.DataFrame(self.progress, columns=['loss'])
        df.plot(figsize=(16, 8), alpha=0.1, marker='.', grid=True)
#%%
class G(nn.Module):
    def __init__(self):
        super(G, self).__init__()
        self.model = nn.Sequential(
            nn.Linear(100 + 10, 200),
            nn.LeakyReLU(0.02),
            nn.LayerNorm(200),
            nn.Linear(200, 784),
            nn.Tanh()
        )

        self.optimiser = torch.optim.Adam(self.parameters(), lr=0.001)
        self.counter = 0
        self.progress = []

    def forward(self, inputs, label_tensors):
        return self.model(torch.cat((inputs, label_tensors), 1))

    def nn_train(self, discriminator: D, inputs, label_tensors, targets):
        g_outputs = self(inputs, label_tensors)
        d_outputs = discriminator(g_outputs, label_tensors)
        loss = discriminator.loss_func(d_outputs, targets)
        self.counter += 1
        if not self.counter % 10:
            self.progress.append(loss.item())
        self.optimiser.zero_grad()
        loss.backward()
        self.optimiser.step()

    def plot_progress(self):
        df = pandas.DataFrame(self.progress, columns=['loss'])
        df.plot(figsize=(16, 8), alpha=0.1, marker='.', grid=True)

    def plot_images(self, label_tensor):
        _, canvases = plt.subplots(1, 4)
        for i in range(4):
            canvases[i].imshow(self(torch.randn(1, 100, device=device), to_hot(torch.tensor([label_tensor]))).cpu().detach().numpy().reshape(28, 28), interpolation='none', cmap='Blues')
#%%
def random_hot(batch_size):
    tensor = torch.zeros(batch_size, 10, device=device)
    random_index = [random.randint(0, 9) for _ in range(batch_size)]
    for t, index in enumerate(random_index):
        tensor[t][index] = 1
    return tensor

def to_hot(labels):
    tensor = torch.zeros(labels.size(0), 10, device=device)
    for t, index in enumerate(labels):
        tensor[t][index] = 1
    return tensor
#%%
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize(0.5, 0.5)])
trainset = datasets.MNIST(root='../data', train=True, download=True, transform=transform)
trainloader = DataLoader(trainset, batch_size=64, shuffle=True)
PATH_D = '../pth/GAN_mnistD.pth'
PATH_G = '../pth/GAN_mnistG.pth'
#%%
d = D().to(device)
g = G().to(device)
#%%
# try:
#     d = torch.load(PATH_D, map_location=device)
#     g = torch.load(PATH_G, map_location=device)
# except FileNotFoundError:
#     d = D().to(device)
#     g = G().to(device)
#%%
epoch = 5
#%%
%%time
for e in range(epoch):
    print(f'epcoh-{e + 1}...', end='')
    for step, (data, label) in enumerate(trainloader):
        data = data.to(device)
        size = data.size(0)
        d.nn_train(data * -1, to_hot(label), torch.ones(size, 1, device=device))
        random_label = random_hot(size)
        d.nn_train(g(torch.randn(size, 100, device=device), random_label), random_label, torch.zeros(size, 1, device=device))
        random_label = random_hot(size)
        g.nn_train(d, torch.randn(size, 100, device=device), random_label, torch.ones(size, 1, device=device))
    print('finished')
#%%
d.plot_progress()
g.plot_progress()
#%%
for n in range(10):
    g.plot_images(n)

判别器loss

生成器loss

生成的图像

就是说效果看着还可以,有的长得比较残但是能看出来是什么数字(想当初模式坍塌真痛苦)

当然现在把这个改成DCGAN还没成功,继续加油吧(等考完期末