ds证据理论python实现_WGANdiv模型理论以及Python实现

https://arxiv.org/pdf/1712.01026.pdf

WGAN-div很明显是基于WGAN的一个很重要的创新。

WGAN在理论上确实发现了一个GAN存在的收敛性上的问题。在WGAN中，WGAN作者通过提出JS测度(GAN)，Wasserstein测度(WGAN)等多种测度的角度出发，去拟合平行线(0, Z)其中Z是[0,1]上的分布。这时候，选取生成器为g_theta(z) = (theta, z)。这时候，就会发现衡量生成器的分布和目标分布(0, Z)直接的测度关系，有：

会发现，这里，只有W测度(EM距离)，对应的是连续的，其他却不能保证连续。

换个角度看，这里的theta对应对于模型本身的分布的参数(也就是神经网络G的参数)，在输入的z之后的输出分布和原来的真实部分之间是平行关系的时候，其他的测度比如(JS)对应的梯度(相对于模型参数theta)就消失了。

这也就解释了为什么GAN(对应于JS散度)，很容易出现梯度消失的问题。

也就会导致模型输出的结果不一定和真实分布的图像很接近。

比如输出一个人，人分为头和脚(假设分布下两者被生成概率相同)， 那么完全可以出现将头和脚的位置调换的情况。从测度上看是没有问题的，但是从输出的图像上看就不是我们需要的结果。(也一定程度上解释了GAN的效果不太行的原因)。

这就是为什么大家后来都吹WGAN的原因。毕竟从理论上解释了GAN的收敛不稳定问题。但是Wasserstein测度不好求，于是WGAN就提出处理对应问题的对偶问题，但是却引入了lipschitz约束。

这个约束如何在模型的设计中表现出现，就是不同WGAN的设计问题了。

WGAN给出了weight clipping的操作但是不理想。

WGAN-gp给出了gradient penalty的操作。可以说是开了先河(后面的很多WGAN-family模型都是在这个gp上做文章)

但是这一堆WGAN-family的模型中，WGAN-div可以说是集大成者...吧？(其他的我还没实现，反正论文作者把其他都拿出来做对比(论文中显示效果确实更优)，可能是理论上有优势吧，后续再看)

WGAN-div重要的地方就是把这个lipschitz约束给去掉了。

讲到这个改进，就需要设计到了一篇97年发的数学论文(55页)(太长， 之后有时间再看看)。emmm做ai都需要去考古以前的数学论文了，感觉宝藏还不少。

Evans, L.C.: Partial differential equations and monge-kantorovich mass transfer. Current developments in mathematics 1997(1) (1997) 65–126 

也就是上面的这个函数，当p这个指数趋于无穷的时候，这个f的对应就是wgan的Wasserstein测度的最优解的负数上。(那这样一样是完成了求解嘛)

而且这个函数是没有lipschitz约束在的，但是仍然存在p是需要趋于无穷的时候，才能保证趋于WGAN给定EM距离的最小的情况。

但是，根据极限的定义，学过数分的都写下面这段话，对任意常数C，存在某个常数G，s.t. 任意k大于G都有f(k) < C。

所以说，我们只需要找个有点大的p就可能够用了。

因此，作者需要做一堆关于p的测试。如下(第二行左图)：

因为看到p取6的时候效果就还不错了，所以WGAN-div推荐的超参p取6，同理k取2。

---

恰饭

---

实验

做了几组实验测试：

1. 线性模型+wgan-div

2. 线性模型+wgan-div+sigmoid

3. conv测试

4. conv再测

首先，我是先用DCGAN的直接去改的wgan-div。因为之前的WGAN和WGAN-gp就是这么成功。

但是我试了很多次之后，发现wgan一直失败，我就在网上找到成功的代码。发现他们都是使用线性模型的，然后就有了实验1。

用他们的模型，成功之后，发现他们居然没给D最后一层放sigmoid。感觉很无语啊。因此就有了实验2。

但我很不服气，如果改成dcgan就不成功了，那么wgan-div感觉也没啥用吧？于是就再试着探索了下为什么加了卷积层不成功。这就有了实验3。

从实验的结果来看，是D中不应该使用BN层。反正加了BN层的之后，模型就崩溃了。(D超早就直接拒绝G(这？？？不是 梯度消失导致的嘛，wgan不是就是解决这个的嘛，怎么被这玩意搞定了。))

最后再验证是否是BN存在的问题，就把DCGAN中的D的BN删掉之后，再做测试。也就是实验4。

个人猜测(不一定准确)：可能还是梯度的问题。本质上还是有没有log嘛。log虽然在接近1的时候， 会很缓。但是在接近0的时候就会相当陡。但是换成wgan这样之后，log被去掉了。那么这么来就导致一开始的时候(数值接近0的时候)，梯度太小，训练结果不好。删掉BN之后，使得D进步的慢一点，就不那么容易一开始就拒绝G。(G的零启动问题)

效果提前展示，代码统一放在了最后。

实验一效果：

线性模型+wgan-div

实验二效果：

线性模型+wgan-div+sigmoid

实验三效果：

conv测试

实验4效果：

conv再测

实验中线性的部分都是用epoch 50，卷积的部分使用的是epoch20。在这种情况下，实验4仍然是较好的结果。

实验1代码：

main.py

import osimport torchfrom torch.utils.data import Dataset, DataLoaderfrom model import Generator, Discriminator, div_lossimport torchvisionimport matplotlib.pyplot as pltif __name__ == '__main__':    LR = 0.0002    EPOCH = 50  # 50    BATCH_SIZE = 100    N_IDEAS = 100    nc = 2    TRAINED = False    k = 2    p = 6    DOWNLOAD_MNIST = False    mnist_root = '../Conditional-GAN/mnist/'    if not (os.path.exists(mnist_root)) or not os.listdir(mnist_root):        # not mnist dir or mnist is empyt dir        DOWNLOAD_MNIST = True    train_data = torchvision.datasets.MNIST(        root=mnist_root,        train=True,  # this is training data        transform=torchvision.transforms.ToTensor(),  # Converts a PIL.Image or numpy.ndarray to        # torch.FloatTensor of shape (C x H x W) and normalize in the range [0.0, 1.0]        download=DOWNLOAD_MNIST,    )    train_loader = DataLoader(dataset=train_data, batch_size=BATCH_SIZE, shuffle=True)    torch.cuda.empty_cache()    if TRAINED:        G = torch.load('G.pkl').cuda()        D = torch.load('D.pkl').cuda()    else:        G = Generator(N_IDEAS).cuda()        D = Discriminator().cuda()    optimizerG = torch.optim.Adam(G.parameters(), lr=LR)    optimizerD = torch.optim.Adam(D.parameters(), lr=LR)    for epoch in range(EPOCH):        tmpD, tmpG = 0, 0        for step, (x, y) in enumerate(train_loader):            x = x.cuda()            rand_noise = torch.randn((x.shape[0], N_IDEAS)).cuda()            G_imgs = G(rand_noise)            D_fake = torch.squeeze(D(G_imgs))            D_real = torch.squeeze(D(x))            D_real_loss = -torch.mean(D_real)            D_fake_loss = torch.mean(D_fake)            D_loss = D_real_loss + D_fake_loss + k * div_loss(D, x, G_imgs, p=p, cuda=True)            optimizerD.zero_grad()            D_loss.backward(retain_graph=True)            optimizerD.step()            if (step + 1) % nc == 0:                G_loss = -torch.mean(D_fake)                optimizerG.zero_grad()                G_loss.backward(retain_graph=True)                optimizerG.step()                tmpG_ = G_loss.cpu().detach().data                tmpG += tmpG_                tmpD_ = D_loss.cpu().detach().data                tmpD += tmpD_        tmpD /= (step + 1)        tmpG /= (step + 1)        print(            'epoch %d avg of loss: D: %.6f, G: %.6f' % (epoch, tmpD * nc, tmpG * nc)        )        if epoch % 2 == 0:            plt.imshow(torch.squeeze(G_imgs[0].cpu().detach()).numpy())            plt.show()    torch.save(G, 'G.pkl')    torch.save(D, 'D.pkl')

model.py

import osimport torchimport torch.nn as nnimport torch.utils.data as Dataimport torchvisionimport torch.autograd as autogradimport numpy as npclass Generator(nn.Module):    def __init__(self, latent_dim):        super(Generator, self).__init__()        self.img_shape = (1, 28, 28)        def block(in_feat, out_feat, normalize=True):            layers = [nn.Linear(in_feat, out_feat)]            if normalize:                layers.append(nn.BatchNorm1d(out_feat, 0.8))            layers.append(nn.LeakyReLU(0.2, inplace=True))            return layers        self.model = nn.Sequential(            *block(latent_dim, 128, normalize=False),            *block(128, 256),            *block(256, 512),            *block(512, 1024),            nn.Linear(1024, int(np.prod(self.img_shape))),            nn.Tanh()        )    def forward(self, z):        img = self.model(z)        img = img.view(img.shape[0], *self.img_shape)        return imgclass Discriminator(nn.Module):    def __init__(self):        super(Discriminator, self).__init__()        self.img_shape = (1, 28, 28)        self.model = nn.Sequential(            nn.Linear(int(np.prod(self.img_shape)), 512),            nn.LeakyReLU(0.2, inplace=True),            nn.Linear(512, 256),            nn.LeakyReLU(0.2, inplace=True),            nn.Linear(256, 1),        )    def forward(self, img):        img_flat = img.view(img.shape[0], -1)        validity = self.model(img_flat)        return validitydef div_loss(D, real_x, fake_x, p=6, cuda=False):    if cuda:        alpha = torch.rand((real_x.shape[0], 1, 1, 1)).cuda()    else:        alpha = torch.rand((real_x.shape[0], 1, 1, 1))    x_ = (alpha * real_x + (1-alpha) * fake_x).requires_grad_(True)    y_ = D(x_)    # cal f'(x)    grad = autograd.grad(        outputs=y_,        inputs=x_,        grad_outputs=torch.ones_like(y_),        create_graph=True,        retain_graph=True,        only_inputs=True,    )[0]    grad = grad.view(x_.shape[0], -1)    div = (grad.norm(2, dim=1) ** p).mean()    return divif __name__ == '__main__':    N_IDEAS = 100    G = Generator(N_IDEAS, )    # rand_noise = torch.randn((10, N_IDEAS, 1, 1))    # print(G(rand_noise).shape)    DOWNLOAD_MNIST = False    mnist_root = '../Conditional-GAN/mnist/'    if not (os.path.exists(mnist_root)) or not os.listdir(mnist_root):        # not mnist dir or mnist is empyt dir        DOWNLOAD_MNIST = True    train_data = torchvision.datasets.MNIST(        root=mnist_root,        train=True,  # this is training data        transform=torchvision.transforms.ToTensor(),  # Converts a PIL.Image or numpy.ndarray to        # torch.FloatTensor of shape (C x H x W) and normalize in the range [0.0, 1.0]        download=DOWNLOAD_MNIST,    )    D = Discriminator()    print(len(train_data))    train_loader = Data.DataLoader(dataset=train_data, batch_size=2, shuffle=True)    for step, (x, y) in enumerate(train_loader):        print(x.shape)        print(D(x).shape)        rand_noise = torch.randn((x.shape[0], N_IDEAS))  # .cuda()        G_imgs = G(rand_noise)        print(div_loss(D, x, G_imgs))        break

judge.py

import numpy as npimport torchimport matplotlib.pyplot as pltfrom model import Generator, Discriminatorimport torchvision.utils as vutilsif __name__ == '__main__':    BATCH_SIZE = 100    N_IDEAS = 100    img_shape = (1, 28, 28)    TIME = 5    G = torch.load("G.pkl").cuda()    for t in range(TIME):        rand_noise = torch.randn((BATCH_SIZE, N_IDEAS)).cuda()        G_imgs = G(rand_noise).cpu().detach()        fig = plt.figure(figsize=(10, 10))        plt.axis("off")        plt.imshow(np.transpose(vutils.make_grid(G_imgs, nrow=10, padding=0, normalize=True, scale_each=True), (1, 2, 0)))        plt.savefig(str(t) + '.png')        plt.show()

实验2代码

相比于实验一只是改了下model.py

import osimport torchimport torch.nn as nnimport torch.utils.data as Dataimport torchvisionimport torch.autograd as autogradimport numpy as npclass Generator(nn.Module):    def __init__(self, latent_dim):        super(Generator, self).__init__()        self.img_shape = (1, 28, 28)        def block(in_feat, out_feat, normalize=True):            layers = [nn.Linear(in_feat, out_feat)]            if normalize:                layers.append(nn.BatchNorm1d(out_feat, 0.8))            layers.append(nn.LeakyReLU(0.2, inplace=True))            return layers        self.model = nn.Sequential(            *block(latent_dim, 128, normalize=False),            *block(128, 256),            *block(256, 512),            *block(512, 1024),            nn.Linear(1024, int(np.prod(self.img_shape))),            nn.Tanh()        )    def forward(self, z):        img = self.model(z)        img = img.view(img.shape[0], *self.img_shape)        return imgclass Discriminator(nn.Module):    def __init__(self):        super(Discriminator, self).__init__()        self.img_shape = (1, 28, 28)        self.model = nn.Sequential(            nn.Linear(int(np.prod(self.img_shape)), 512),            nn.LeakyReLU(0.2, inplace=True),            nn.Linear(512, 256),            nn.LeakyReLU(0.2, inplace=True),            nn.Linear(256, 1),            nn.Sigmoid()        )    def forward(self, img):        img_flat = img.view(img.shape[0], -1)        validity = self.model(img_flat)        return validitydef div_loss(D, real_x, fake_x, p=6, cuda=False):    if cuda:        alpha = torch.rand((real_x.shape[0], 1, 1, 1)).cuda()    else:        alpha = torch.rand((real_x.shape[0], 1, 1, 1))    x_ = (alpha * real_x + (1-alpha) * fake_x).requires_grad_(True)    y_ = D(x_)    # cal f'(x)    grad = autograd.grad(        outputs=y_,        inputs=x_,        grad_outputs=torch.ones_like(y_),        create_graph=True,        retain_graph=True,        only_inputs=True,    )[0]    grad = grad.view(x_.shape[0], -1)    div = (grad.norm(2, dim=1) ** p).mean()    return divif __name__ == '__main__':    N_IDEAS = 100    G = Generator(N_IDEAS, )    # rand_noise = torch.randn((10, N_IDEAS, 1, 1))    # print(G(rand_noise).shape)    DOWNLOAD_MNIST = False    mnist_root = '../../Conditional-GAN/mnist/'    if not (os.path.exists(mnist_root)) or not os.listdir(mnist_root):        # not mnist dir or mnist is empyt dir        DOWNLOAD_MNIST = True    train_data = torchvision.datasets.MNIST(        root=mnist_root,        train=True,  # this is training data        transform=torchvision.transforms.ToTensor(),  # Converts a PIL.Image or numpy.ndarray to        # torch.FloatTensor of shape (C x H x W) and normalize in the range [0.0, 1.0]        download=DOWNLOAD_MNIST,    )    D = Discriminator()    print(len(train_data))    train_loader = Data.DataLoader(dataset=train_data, batch_size=2, shuffle=True)    for step, (x, y) in enumerate(train_loader):        print(x.shape)        print(D(x).shape)        rand_noise = torch.randn((x.shape[0], N_IDEAS))  # .cuda()        G_imgs = G(rand_noise)        print(div_loss(D, x, G_imgs))        break

实验3代码

因为G这时候前面是加了线性映射层的，所以就不用改其他代码，只需要修改model.py就好了。

import osimport torchimport torch.nn as nnimport torch.utils.data as Dataimport torchvisionimport torch.autograd as autogradimport numpy as npdef conv_T_block(in_feat, out_feat, normalize=True):    layers = [nn.ConvTranspose2d(in_feat, out_feat, kernel_size=4, stride=2, padding=1)]    if normalize:        layers.append(nn.BatchNorm2d(out_feat, 0.8))    layers.append(nn.LeakyReLU(0.2, inplace=True))    return layersdef conv_block(in_feat, out_feat, normalize=True):    layers = [nn.Conv2d(in_feat, out_feat, kernel_size=4, stride=2, padding=1)]    if normalize:        layers.append(nn.BatchNorm2d(out_feat, 0.8))    layers.append(nn.LeakyReLU(0.2, inplace=True))    return layersdef block(in_feat, out_feat, normalize=True):    layers = [nn.Linear(in_feat, out_feat)]    if normalize:        layers.append(nn.BatchNorm1d(out_feat, 0.8))    layers.append(nn.LeakyReLU(0.2, inplace=True))    return layersclass Generator(nn.Module):    def __init__(self, latent_dim):        super(Generator, self).__init__()        self.img_shape = (1, 28, 28)        self.init_img_shape = (64, 7, 7)        self.fc = nn.Sequential(            *block(latent_dim, np.prod(self.init_img_shape), normalize=False),        )        self.model = nn.Sequential(            *conv_T_block(self.init_img_shape[0], self.init_img_shape[0] // 2),            *conv_T_block(self.init_img_shape[0] // 2, 1),            nn.Tanh()        )    def forward(self, z):        img = self.fc(z).view(z.shape[0], *self.init_img_shape)        img = self.model(img)        return imgclass Discriminator(nn.Module):    def __init__(self):        super(Discriminator, self).__init__()        self.img_shape = (1, 28, 28)        self.final_img_shape = (64, 7, 7)        self.model = nn.Sequential(            *conv_block(self.img_shape[0], self.final_img_shape[0] // 2, normalize=False),            *conv_block(self.final_img_shape[0] // 2, self.final_img_shape[0], normalize=False)        )        self.fc = nn.Sequential(            nn.Linear(int(np.prod(self.final_img_shape)), 256),            nn.LeakyReLU(0.2, inplace=True),            nn.Linear(256, 1),            nn.Sigmoid()        )    def forward(self, img):        img_conv = self.model(img)        img_flat = img_conv.view(img.shape[0], -1)        validity = self.fc(img_flat)        return validitydef div_loss(D, real_x, fake_x, p=6, cuda=False):    if cuda:        alpha = torch.rand((real_x.shape[0], 1, 1, 1)).cuda()    else:        alpha = torch.rand((real_x.shape[0], 1, 1, 1))    x_ = (alpha * real_x + (1 - alpha) * fake_x).requires_grad_(True)    y_ = D(x_)    # cal f'(x)    grad = autograd.grad(        outputs=y_,        inputs=x_,        grad_outputs=torch.ones_like(y_),        create_graph=True,        retain_graph=True,        only_inputs=True,    )[0]    grad = grad.view(x_.shape[0], -1)    div = (grad.norm(2, dim=1) ** p).mean()    return divif __name__ == '__main__':    N_IDEAS = 100    G = Generator(N_IDEAS)    rand_noise = torch.randn((10, N_IDEAS))    print(G(rand_noise).shape)    DOWNLOAD_MNIST = False    mnist_root = '../../Conditional-GAN/mnist/'    if not (os.path.exists(mnist_root)) or not os.listdir(mnist_root):        # not mnist dir or mnist is empyt dir        DOWNLOAD_MNIST = True    train_data = torchvision.datasets.MNIST(        root=mnist_root,        train=True,  # this is training data        transform=torchvision.transforms.ToTensor(),  # Converts a PIL.Image or numpy.ndarray to        # torch.FloatTensor of shape (C x H x W) and normalize in the range [0.0, 1.0]        download=DOWNLOAD_MNIST,    )    D = Discriminator()    print(len(train_data))    train_loader = Data.DataLoader(dataset=train_data, batch_size=2, shuffle=True)    for step, (x, y) in enumerate(train_loader):        print(x.shape)        print(D(x).shape)        rand_noise = torch.randn((x.shape[0], N_IDEAS))  # .cuda()        G_imgs = G(rand_noise)        print(div_loss(D, x, G_imgs))        break

实验4代码

相比起来，因为需要同时修改输入给G的z的规模并且最后的结果也是最优的，因此这里就放一下完整代码。

main.py

import osimport torchfrom torch.utils.data import Dataset, DataLoaderfrom model import Generator, Discriminator, div_lossimport torchvisionimport matplotlib.pyplot as pltif __name__ == '__main__':    LR = 0.0002    EPOCH = 20  # 50    BATCH_SIZE = 100    N_IDEAS = 100    nc = 2    TRAINED = False    k = 2    p = 6    DOWNLOAD_MNIST = False    mnist_root = '../../Conditional-GAN/mnist/'    if not (os.path.exists(mnist_root)) or not os.listdir(mnist_root):        # not mnist dir or mnist is empyt dir        DOWNLOAD_MNIST = True    train_data = torchvision.datasets.MNIST(        root=mnist_root,        train=True,  # this is training data        transform=torchvision.transforms.ToTensor(),  # Converts a PIL.Image or numpy.ndarray to        # torch.FloatTensor of shape (C x H x W) and normalize in the range [0.0, 1.0]        download=DOWNLOAD_MNIST,    )    train_loader = DataLoader(dataset=train_data, batch_size=BATCH_SIZE, shuffle=True)    torch.cuda.empty_cache()    if TRAINED:        G = torch.load('G.pkl').cuda()        D = torch.load('D.pkl').cuda()    else:        G = Generator(N_IDEAS).cuda()        D = Discriminator().cuda()    optimizerG = torch.optim.Adam(G.parameters(), lr=LR)    optimizerD = torch.optim.Adam(D.parameters(), lr=LR)    for epoch in range(EPOCH):        tmpD, tmpG = 0, 0        for step, (x, y) in enumerate(train_loader):            x = x.cuda()            rand_noise = torch.randn((x.shape[0], N_IDEAS, 1, 1)).cuda()            G_imgs = G(rand_noise)            D_fake = torch.squeeze(D(G_imgs))            D_real = torch.squeeze(D(x))            D_real_loss = -torch.mean(D_real)            D_fake_loss = torch.mean(D_fake)            D_loss = D_real_loss + D_fake_loss + k * div_loss(D, x, G_imgs, p=p, cuda=True)            optimizerD.zero_grad()            D_loss.backward(retain_graph=True)            optimizerD.step()            if (step + 1) % nc == 0:                G_loss = -torch.mean(D_fake)                optimizerG.zero_grad()                G_loss.backward(retain_graph=True)                optimizerG.step()                tmpG_ = G_loss.cpu().detach().data                tmpG += tmpG_                tmpD_ = D_loss.cpu().detach().data                tmpD += tmpD_        tmpD /= (step + 1)        tmpG /= (step + 1)        print(            'epoch %d avg of loss: D: %.6f, G: %.6f' % (epoch, tmpD * nc, tmpG * nc)        )        if epoch % 2 == 0:            plt.imshow(torch.squeeze(G_imgs[0].cpu().detach()).numpy())            plt.show()    torch.save(G, 'G.pkl')    torch.save(D, 'D.pkl')

model.py

import osimport torch.autograd as autogradimport torchimport torch.nn as nnimport torch.utils.data as Dataimport torchvisionfrom torch.utils.data import DataLoaderclass Generator(nn.Module):    def __init__(self, input_size):        super(Generator, self).__init__()        strides = [1, 2, 2, 2]        padding = [0, 1, 1, 1]        channels = [input_size,                    256, 128, 64, 1]  # 1表示一维        kernels = [4, 3, 4, 4]        model = []        for i, stride in enumerate(strides):            model.append(                nn.ConvTranspose2d(                    in_channels=channels[i],                    out_channels=channels[i + 1],                    stride=stride,                    kernel_size=kernels[i],                    padding=padding[i]                )            )            if i != len(strides) - 1:                model.append(                    nn.BatchNorm2d(channels[i + 1], 0.8)                )                model.append(                    nn.LeakyReLU(.2)                )            else:                model.append(                    nn.Tanh()                )        self.main = nn.Sequential(*model)    def forward(self, x):        x = self.main(x)        return xclass Discriminator(nn.Module):    def __init__(self, input_size=1):        super(Discriminator, self).__init__()        strides = [2, 2, 2, 1]        padding = [1, 1, 1, 0]        channels = [input_size,                    64, 128, 256, 1]  # 1表示一维        kernels = [4, 4, 4, 3]        model = []        for i, stride in enumerate(strides):            model.append(                nn.Conv2d(                    in_channels=channels[i],                    out_channels=channels[i + 1],                    stride=stride,                    kernel_size=kernels[i],                    padding=padding[i]                )            )            if i != len(strides) - 1:                model.append(                    nn.LeakyReLU(0.2)                )            else:                model.append(                    nn.Sigmoid()                )        self.main = nn.Sequential(*model)    def forward(self, x):        x = self.main(x)        return xdef div_loss(D, real_x, fake_x, p=6, cuda=False):    if cuda:        alpha = torch.rand((real_x.shape[0], 1, 1, 1)).cuda()    else:        alpha = torch.rand((real_x.shape[0], 1, 1, 1))    x_ = (alpha * real_x + (1 - alpha) * fake_x).requires_grad_(True)    y_ = D(x_)    # cal f'(x)    grad = autograd.grad(        outputs=y_,        inputs=x_,        grad_outputs=torch.ones_like(y_),        create_graph=True,        retain_graph=True,        only_inputs=True,    )[0]    grad = grad.view(x_.shape[0], -1)    div = (grad.norm(2, dim=1) ** p).mean()    return divif __name__ == '__main__':    N_IDEAS = 100    G = Generator(N_IDEAS, )    rand_noise = torch.randn((10, N_IDEAS, 1, 1))    print(G(rand_noise).shape)    DOWNLOAD_MNIST = False    mnist_root = '../../Conditional-GAN/mnist/'    if not (os.path.exists(mnist_root)) or not os.listdir(mnist_root):        # not mnist dir or mnist is empyt dir        DOWNLOAD_MNIST = True    train_data = torchvision.datasets.MNIST(        root=mnist_root,        train=True,  # this is training data        transform=torchvision.transforms.ToTensor(),  # Converts a PIL.Image or numpy.ndarray to        # torch.FloatTensor of shape (C x H x W) and normalize in the range [0.0, 1.0]        download=DOWNLOAD_MNIST,    )    D = Discriminator(1)    print(len(train_data))    train_loader = Data.DataLoader(dataset=train_data, batch_size=2, shuffle=True)    for step, (x, y) in enumerate(train_loader):        print(x.shape)        print(D(x).shape)        break

judge.py

import numpy as npimport torchimport matplotlib.pyplot as pltfrom model import Generator, Discriminatorimport torchvision.utils as vutilsif __name__ == '__main__':    BATCH_SIZE = 100    N_IDEAS = 100    img_shape = (1, 28, 28)    TIME = 5    G = torch.load("G.pkl").cuda()    for t in range(TIME):        rand_noise = torch.randn((BATCH_SIZE, N_IDEAS, 1, 1)).cuda()        G_imgs = G(rand_noise).cpu().detach()        fig = plt.figure(figsize=(10, 10))        plt.axis("off")        plt.imshow(np.transpose(vutils.make_grid(G_imgs, nrow=10, padding=0, normalize=True, scale_each=True), (1, 2, 0)))        plt.savefig(str(t) + '.png')        plt.show()