WGAN的代码实现

1.哔站视频链接：

https://www.bilibili.com/video/BV1TU4y1H7Mz?spm_id_from=333.1007.top_right_bar_window_custom_collection.content.click

2.github链接：

https://github.com/dragen1860/Deep-Learning-with-PyTorch-Tutorials/blob/master/lesson57-WGAN%E5%AE%9E%E6%88%98/wgan_gp.py

3.我的代码注释：

# 说明如果遇见以下错误TypeError: can't convert cuda:0 device type tensor to numpy. Use Tensor.cpu() to copy the tensor to host memory first.
#  在此文件中File "E:\anaconda\lib\site-packages\torch\_tensor.py", line 678, in __array__（这个文件是你运行后爆红的最后一个错误提示）
#     第678行修改return self.numpy()为return self.cpu().numpy()


import torch
import torch.nn as nn
import numpy as np
import random
import visdom
import matplotlib.pyplot as plt

# 设置设备
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# 超参数
h_dim = 400
batch_size = 512
viz = visdom.Visdom()  # 记得cmd中    python -m visdom.server--》按照地址打开即可显示



# 创建模型结构   生成器  判别器
class Generator(nn.Module):
    def __init__(self):
        super(Generator, self).__init__()
        self.net = nn.Sequential(
            # z:[b,2]--->[b,2]
            # Linear--ReLU--Linear--ReLU--Linear--ReLU--Linear
            nn.Linear(2, h_dim),
            nn.ReLU(True),
            nn.Linear(h_dim, h_dim),
            nn.ReLU(inplace=True),
            nn.Linear(h_dim, h_dim),
            nn.ReLU(inplace=True),
            nn.Linear(h_dim, 2)
        )

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


class Discriminator(nn.Module):
    def __init__(self):
        super(Discriminator, self).__init__()
        self.net = nn.Sequential(
            # z:[b,2]--->[b,1]
            # Linear--ReLU--Linear--ReLU--Linear--ReLU--Linear
            nn.Linear(2, h_dim),
            nn.ReLU(True),
            nn.Linear(h_dim, h_dim),
            nn.ReLU(inplace=True),
            nn.Linear(h_dim, h_dim),
            nn.ReLU(inplace=True),
            nn.Linear(h_dim, 1),
            nn.Sigmoid()  # 设置范围为【0,1】表示当前输入是真实分布的程度
        )

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


# 生成数据集
def data_generator():
    '''
    8-gaussian mixture models 高斯混合模型--》详见mk笔记
    :return:
    '''
    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(batch_size):
            point = np.random.randn(2) * 0.02  # point的shape为（1*2）
            center = random.choice(centers)  # 从给定的八个点中随机选一个点

            # 将point与center对应的位置相加组成新的一个点 加到数据集中
            point[0] += center[0]
            point[1] += center[1]
            dataset.append(point)

        #此时dataset中的元素为一个个的array，经过array成为一个matrix
        dataset = np.array(dataset).astype(np.float32)
        dataset /= 1.414
        yield dataset   #yield是一个生成器


# 生成图像
def generate_image(D, G, x_r, epoch):
    """
    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)
    # print('p:', points.shape)

    # draw contour
    with torch.no_grad():
        points = torch.Tensor(points).cuda()  # [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(batch_size, 2).cuda()  # [b, 2]
        samples = G(z).cpu().numpy()  # [b, 2]
    plt.scatter(x_r[:, 0], x_r[:, 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))


#NN的参数模型初始化
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)

def gradient_penalty(D,x_r,x_f):
    '''

    :param D:
    :param x_r: [b,2]
    :param x_f: [b,2]
    :return:
    '''

  #根据x的线性插值的公式
    #维度：[b,1]
    t=torch.rand(batch_size,1).to(device)  #t[0,1]
    #扩展t的维度与x_r相同，为[b,2]
    t.expand_as(x_r)

    #interpolation(插值)
    mid=t*x_r+(1-t)*x_f
    #设置mid可以求导  （因为公式中GP部分需要导数）
    mid.requires_grad_()

    #得到输出
    pred=D(mid)
    #求导数
    grads=torch.autograd.grad(outputs=pred,
                         inputs=mid,

                         grad_outputs=torch.ones_like(pred),
                         create_graph=True,#这是为了二阶求导
                         retain_graph=True,#如果需要backward则需要保留梯度信息
                         only_inputs=True
                         )[0]
    # grads.norm(2,dim=1)  #2范数
    #根据Loss的计算公式可知  2范数与1越接近越好，
    gp=torch.pow(grads.norm(2, dim=1)-1,2).mean()
    return gp






def main():
    torch.manual_seed(23)
    np.random.seed(23)

    # 实例化G与D
    # 0.3版本
    # G=Generator().cuda()
    # D=Discriminator().cuda()
    # 0.4版本
    G = Generator().to(device)
    D = Discriminator().to(device)

    G.apply(weights_init)
    D.apply(weights_init)

    # 生成数据
    data_iter = data_generator()
    # x=next(data_iter)
    # print(x.shape)     #（512,2）

    # print(G) #打印网络结构
    # print(D)
    # 设置G与D优化器
    optim_G = torch.optim.Adam(G.parameters(), lr=1e-3, betas=(0.5, 0.9))
    optim_D = torch.optim.Adam(D.parameters(), lr=1e-3, betas=(0.5, 0.9))

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

    # GAN的核心部分
    for epoch in range(5000):

        # 1.train Discriminator Firstly
        for _ in range(5):  # 暂定优化五步
            # ①训练real_data
            x_r = next(data_iter) #一组为batch_size大小
            # 数据处理：把数据设置为tensor
            x_r = torch.from_numpy(x_r).to(device)
            # [d,2]-->[d,1]
            pred_r = (D(x_r))
            # 最大化pred_r,最小化loss_r(所以加上符号)
            loss_r = -(pred_r.mean())

            # ②训练fake_data
            # 创建假数据
            z = torch.randn(batch_size, 2).to(device)
            x_f = G(z).detach()  # 因为不优化G，detach()梯度计算停止

            # 判别假数据
            pred_f = (D(x_f))
            loss_f = (pred_f.mean())

            #设置GP
            gp=gradient_penalty(D,x_r,x_f.detach())  #因为是优化D,不需要对G求导，所以要对x_f进行detach()

            # ③合并损失
            loss_D = loss_f + loss_r+0.2*gp

            # ④优化
            optim_D.zero_grad()
            loss_D.backward()
            optim_D.step()

        # 2.train Generator
        # 创建fake数据--》初始给G的噪声一样
        z2 = torch.randn(batch_size, 2).to(device)
        # 使用G生
        x_f2 = G(z2)
        pred_f2 = (D(x_f2))  # 这里虽然不优化D但是也没有用detach因为他在G的后面
        # 这是优化G，所以需要pred_f2越大越好，loss_G越小越好（故取负数）
        loss_G = -(pred_f2.mean())

        # 优化
        optim_G.zero_grad()
        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, x_r, epoch)
            print(loss_D.item(), loss_G.item())


if __name__ == "__main__":
    main()

4.补充

4.1yield的作用

https://blog.csdn.net/mieleizhi0522/article/details/82142856/