Pytorch-GAN

任务：使用8个高斯混合模型生成一系列数据，通过GAN学习它的分布，比较学习的分布和真实的分布是否一样。

GAN文字版算法：

GAN公式版算法：

在命令行执行如下语句（详细Visdom的使用见https://www.cnblogs.com/cxq1126/p/13285150.html）

python -m visidom.server

提前导包：

import  torch 
from    torch import nn, optim, autograd
import  numpy as np
import  visdom
from    torch.nn import functional as F
from    matplotlib import pyplot as plt
import  random

h_dim = 400
batchsz = 512
viz = visdom.Visdom()

1.实现Generator

输入点坐标，输出点坐标。

class Generator(nn.Module):

    def __init__(self):
        super(Generator, self).__init__()

        self.net = nn.Sequential(             
            nn.Linear(2, h_dim),          
            nn.ReLU(True),
            nn.Linear(h_dim, h_dim),
            nn.ReLU(True),
            nn.Linear(h_dim, h_dim),
            nn.ReLU(True),
            nn.Linear(h_dim, 2),
        )

    def forward(self, z):
        output = self.net(z)
        return output

2.实现Discriminator

输入点坐标，输出数值（用来评判输入的坐标是否在真实数据附近）。

class Discriminator(nn.Module):

    def __init__(self):
        super(Discriminator, self).__init__()

        self.net = nn.Sequential(
            nn.Linear(2, h_dim),
            nn.ReLU(True),
            nn.Linear(h_dim, h_dim),
            nn.ReLU(True),
            nn.Linear(h_dim, h_dim),
            nn.ReLU(True),
            nn.Linear(h_dim, 1),
            nn.Sigmoid()
        )

    def forward(self, x):
        output = self.net(x)
        return output.view(-1)

3.权重初始化

def weights_init(m):
    if isinstance(m, nn.Linear):
        # m.weight.data.normal_(0.0, 0.02)
        nn.init.kaiming_normal_(m.weight)
        m.bias.data.fill_(0)

4.生成数据集 8-gaussian mixture models

对于高斯混合模型的理解：

def data_generator():

    scale = 2.
    centers = [
        (1, 0),
        (-1, 0),
        (0, 1),
        (0, -1),
        (1. / np.sqrt(2), 1. / np.sqrt(2)),
        (1. / np.sqrt(2), -1. / np.sqrt(2)),
        (-1. / np.sqrt(2), 1. / np.sqrt(2)),
        (-1. / np.sqrt(2), -1. / np.sqrt(2))
    ]
    centers = [(scale * x, scale * y) for x, y in centers]
    while True:
        dataset = []
        for i in range(batchsz):
            point = np.random.randn(2) * .02
            center = random.choice(centers)
            
            #N(0,1)sample出来一个点 + center_x1/x2
            point[0] += center[0]
            point[1] += center[1]
            dataset.append(point)
        dataset = np.array(dataset, dtype='float32')
        dataset /= 1.414                                #stdev
        yield dataset

5.可视化

def generate_image(D, G, xr, epoch):      #xr表示真实的sample
    """
    Generates and saves a plot of the true distribution, the generator, and the
    critic.
    """
    N_POINTS = 128
    RANGE = 3
    plt.clf()

    points = np.zeros((N_POINTS, N_POINTS, 2), dtype='float32')
    points[:, :, 0] = np.linspace(-RANGE, RANGE, N_POINTS)[:, None]
    points[:, :, 1] = np.linspace(-RANGE, RANGE, N_POINTS)[None, :]
    points = points.reshape((-1, 2))             # (16384, 2)
    

    # draw contour
    with torch.no_grad():
        points = torch.Tensor(points)      # [16384, 2]
        disc_map = D(points).cpu().numpy() # [16384]
    x = y = np.linspace(-RANGE, RANGE, N_POINTS)
    cs = plt.contour(x, y, disc_map.reshape((len(x), len(y))).transpose())
    plt.clabel(cs, inline=1, fontsize=10)
    # plt.colorbar()


    # draw samples
    with torch.no_grad():
        z = torch.randn(batchsz, 2)                 # [b, 2]
        samples = G(z).cpu().numpy()                # [b, 2]
    plt.scatter(xr[:, 0], xr[:, 1], c='orange', marker='.')
    plt.scatter(samples[:, 0], samples[:, 1], c='green', marker='+')

    viz.matplot(plt, win='contour', opts=dict(title='p(x):%d'%epoch))

6.训练

和上面图片中的梯度上升法不同，下面训练使用的梯度下降法，所以对于原本需要最大化的的数据添加负号，就能实现梯度下降。

optim.Adam()中的batas参数用于计算梯度的平均和平方的系数(默认: (0.9, 0.999))

def main():

    torch.manual_seed(23)
    np.random.seed(23)

    G = Generator()
    D = Discriminator()
    G.apply(weights_init)
    D.apply(weights_init)

    optim_G = optim.Adam(G.parameters(), lr=1e-3, betas=(0.5, 0.9)) 
    optim_D = optim.Adam(D.parameters(), lr=1e-3, betas=(0.5, 0.9))


    data_iter = data_generator()
    print('batch:', next(data_iter).shape)              #[b, 2]

    viz.line([[0,0]], [0], win='loss', opts=dict(title='loss', legend=['D', 'G']))

    for epoch in range(1000):

        # 1. train discriminator for k steps
        for _ in range(5):
            
            #1.1首先train on real data
            x = next(data_iter)
            xr = torch.from_numpy(x)       #真实数据
            predr = (D(xr))                # [b, 2] -> [b, 1]
            # max log(lossr)，即min (-lossr)
            lossr = - (predr.mean())

            #1.2 train on fake data
            z = torch.randn(batchsz, 2)    # [b, 2]随机产生的伪数据
            xf = G(z).detach()             # [b, 2] 此处固定G，更新D，所以不更新G的参数
            predf = (D(xf))                # [b]
            # min predf
            lossf = (predf.mean())

            loss_D = lossr + lossf 
            
            optim_D.zero_grad()
            loss_D.backward()
            optim_D.step()


        # 2. train Generator
        z = torch.randn(batchsz, 2)       #[b, 2]随机产生的伪数据
        xf = G(z)
        predf = (D(xf))
        # max predf，即min(-predf)
        loss_G = - (predf.mean())        
        
        optim_G.zero_grad()               #此处固定D，更新G，所以不更新D的参数
        loss_G.backward()
        optim_G.step()


        if epoch % 100 == 0:
            viz.line([[loss_D.item(), loss_G.item()]], [epoch], win='loss', update='append')
            generate_image(D, G, xr, epoch)
            print(loss_D.item(), loss_G.item())


if __name__ == '__main__':
    main()

下图显示Generator和Discriminator的训练结果，两者的loss都接近于0，sample出来的数据经过Generator后覆盖住了真实数据。

-0.6137181520462036 -0.0665668323636055
7.620922672430937e-23 -6.780519236493468e-22
2.8811570836759978e-37 -3.576674200342663e-41
2.7326036398933606e-12 -5.0371435361452164e-33
1.5811193901845998e-21 -1.796846778152569e-15
2.1587619070328268e-20 -0.0
2.0948376092776535e-32 -2.429269052203856e-16
6.822592214491066e-14 -0.0
9.122023851176224e-35 -2.9085079939065756e-34
0.0 -4.5381907731517276e-14

并不是每次运行都是这种结果，Generator常常由于GAN训练的不稳定（真实数据和生成数据没有重叠），loss保持在非0的某个值，长期得不到更新。

解决方案：W-GAN（通过用Wasserstein距离代替JS散度来优化训练的生成对抗网络）

在代码中增加gradient penalty

def gradient_penalty(D, xr, xf):  #xr和xf的shape=[b, 2]
    
    t = torch.rand(batchsz, 1)   #sample一个均值分布[b, 1]
    t = t.expand_as(xr)          #[b, 1] -> [b, 2]
    
    mid = t *xr +(1-t) * xf      #在真实数据和fake数据之间做线性差值，即图中的xhat
    mid.requires_grad_()         #设置导数信息
    
    pred = D(mid)
    grads = autograd.grad(outputs=pred, inputs=mid,
                          grad_outputs=torch.ones_like(pred),
                          create_graph=True, retain_graph=True, only_inputs=True)[0]     #create_graph用来二次求导
    
    gp = torch.pow(grads.norm(2, dim=1) - 1, 2).mean()
    return gp

在主函数中计算Discriminator的loss时加上gp项。

            #1.3 gradient penalty
            gp = gradient_penalty(D, xr, xf.detach())    #因为不需要对D求导，所以detach
            
            loss_D = lossr + lossf + 0.2 * gp

迭代2000次的结果如下：

-0.5074049830436707 -0.11114723235368729
-0.4972265362739563 -0.3060861825942993
-0.5360251665115356 -0.23698316514492035
-0.3480537533760071 -0.3882521390914917
-0.22527582943439484 -0.5057252645492554
-0.13060471415519714 -0.5396959185600281
-0.07626903057098389 -0.6366142630577087
-0.09713903069496155 -0.6304153203964233
-0.1190759465098381 -0.5412021279335022
-0.1230357214808464 -0.5588557124137878
-0.04560390114784241 -0.6632308959960938
-0.06906679272651672 -0.6173125505447388
-0.04104984924197197 -0.7628952860832214
-0.0408158078789711 -0.7121548652648926
-0.04687119275331497 -0.7424123287200928
-0.024066904559731483 -0.7196884751319885
-0.04576507583260536 -0.7208324670791626
-0.02462894842028618 -0.7012563943862915
-0.01230126153677702 -0.7875514030456543
-0.02122686244547367 -0.7108622193336487