[Pytorchニューラルネットワークの戦闘事例]Fashion-MNSTシミュレーションデータを生成するためのW発散を伴う35WGAN-divモデル

1WGAN-divの概要

W発散GAN-dvモデルの損失関数は、W発散を使用してW距離の計算方法を置き換え、真と偽のサンプルの元のサンプリング操作を分布レベルに基づく計算に置き換えます。

2コードの実装

WGAN-gpに基づいて、わずかな変更を加えて実装され、損失関数の実装が書き直されています。

2.1コードコンバット：モジュールの導入とサンプルのロード----WGAN_div_241.py（パート1）

import torch
import torchvision
from torchvision import transforms
from torch.utils.data import DataLoader
from torch import nn
import torch.autograd as autograd
import matplotlib.pyplot as plt
import numpy as np
import matplotlib
import os
os.environ["KMP_DUPLICATE_LIB_OK"]="TRUE"

# 1.1 引入模块并载入样本:定义基本函数，加载FashionMNIST数据集
def to_img(x):
    x = 0.5 * (x+1)
    x = x.clamp(0,1)
    x = x.view(x.size(0),1,28,28)
    return x

def imshow(img,filename = None):
    npimg = img.numpy()
    plt.axis('off')
    array = np.transpose(npimg,(1,2,0))
    if filename != None:
        matplotlib.image.imsave(filename,array)
    else:
        plt.imshow(array)
        # plt.savefig(filename) # 保存图片 注释掉，因为会报错，暂时不知道什么原因 2022.3.26 15:20
        plt.show()

img_transform = transforms.Compose(
    [
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.5],std=[0.5])
    ]
)

data_dir = './fashion_mnist'

train_dataset = torchvision.datasets.FashionMNIST(data_dir,train=True,transform=img_transform,download=True)
train_loader = DataLoader(train_dataset,batch_size=1024,shuffle=True)
# 测试数据集
val_dataset = torchvision.datasets.FashionMNIST(data_dir,train=False,transform=img_transform)
test_loader = DataLoader(val_dataset,batch_size=10,shuffle=False)
# 指定设备
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
print(device)

2.2コードコンバット：ジェネレーターとディスクリミネーターの実装----WGAN_div_241.py（パート2）

# 1.2 实现生成器和判别器 ：因为复杂部分都放在loss值的计算方面了，所以生成器和判别器就会简单一些。
# 生成器和判别器各自有两个卷积和两个全连接层。生成器最终输出与输入图片相同维度的数据作为模拟样本。
# 判别器的输出不需要有激活函数，并且输出维度为1的数值用来表示结果。
# 在GAN模型中，因判别器的输入则是具体的样本数据，要区分每个数据的分布特征，所以判别器使用实例归一化，
class WGAN_D(nn.Module): # 定义判别器类D ：有两个卷积和两个全连接层
    def __init__(self,inputch=1):
        super(WGAN_D, self).__init__()
        self.conv1 = nn.Sequential(
            nn.Conv2d(inputch,64,4,2,1), # 输出形状为[batch,64,28,28]
            nn.LeakyReLU(0.2,True),
            nn.InstanceNorm2d(64,affine=True)
        )
        self.conv2 = nn.Sequential(
            nn.Conv2d(64,128,4,2,1),# 输出形状为[batch,64,14,14]
            nn.LeakyReLU(0.2,True),
            nn.InstanceNorm2d(128,affine=True)
        )
        self.fc = nn.Sequential(
            nn.Linear(128*7*7,1024),
            nn.LeakyReLU(0.2,True)
        )
        self.fc2 = nn.Sequential(
            nn.InstanceNorm1d(1,affine=True),
            nn.Flatten(),
            nn.Linear(1024,1)
        )
    def forward(self,x,*arg): # 正向传播
        x = self.conv1(x)
        x = self.conv2(x)
        x = x.view(x.size(0),-1)
        x = self.fc(x)
        x = x.reshape(x.size(0),1,-1)
        x = self.fc2(x)
        return x.view(-1,1).squeeze(1)

# 在GAN模型中，因生成器的初始输入是随机值，所以生成器使用批量归一化。
class WGAN_G(nn.Module): # 定义生成器类G：有两个卷积和两个全连接层
    def __init__(self,input_size,input_n=1):
        super(WGAN_G, self).__init__()
        self.fc1 = nn.Sequential(
            nn.Linear(input_size * input_n,1024),
            nn.ReLU(True),
            nn.BatchNorm1d(1024)
        )
        self.fc2 = nn.Sequential(
            nn.Linear(1024,7*7*128),
            nn.ReLU(True),
            nn.BatchNorm1d(7*7*128)
        )
        self.upsample1 = nn.Sequential(
            nn.ConvTranspose2d(128,64,4,2,padding=1,bias=False), # 输出形状为[batch,64,14,14]
            nn.ReLU(True),
            nn.BatchNorm2d(64)
        )
        self.upsample2 = nn.Sequential(
            nn.ConvTranspose2d(64,1,4,2,padding=1,bias=False), # 输出形状为[batch,64,28,28]
            nn.Tanh()
        )
    def forward(self,x,*arg): # 正向传播
        x = self.fc1(x)
        x = self.fc2(x)
        x = x.view(x.size(0),128,7,7)
        x = self.upsample1(x)
        img = self.upsample2(x)
        return img

2.3コードコンバット：w発散を計算します（WGAN-gpベースで新機能）----WGAN_div_241.py（パート3）

# 1.3 计算w散度：返回值充当WGAN-gp中的惩罚项，用于计算判别器的损失
def compute_w_div(real_samples,real_out,fake_samples,fake_out):
    # 定义参数
    k = 2
    p = 6
    # 计算真实空间的梯度
    weight = torch.full((real_samples.size(0),),1,device=device)
    real_grad = autograd.grad(outputs=real_out,
                              inputs=real_samples,
                              grad_outputs=weight,
                              create_graph=True,
                              retain_graph=True,
                              only_inputs=True)[0]
    # L2范数
    real_grad_norm = real_grad.view(real_grad.size(0),-1).pow(2).sum(1)
    # 计算模拟空间的梯度
    fake_grad = autograd.grad(outputs=fake_out,
                              inputs=fake_samples,
                              grad_outputs=weight,
                              create_graph=True,
                              retain_graph=True,
                              only_inputs=True)[0]
    # L2范数
    fake_grad_norm = fake_grad.view(fake_grad.size(0),-1).pow(2).sum(1)
    # 计算W散度距离
    div_gp = torch.mean(real_grad_norm **(p/2)+fake_grad_norm**(p/2))*k/2
    return div_gp

2.4コードコンバット：モデルのトレーニング機能を定義します（WGAN-gp基本変更）----WGAN_div_241.py（パート4）

## 1.4 定义模型的训练函数
# 定义函数train()，实现模型的训练过程。
# 在函数train()中，按照对抗神经网络专题（一）中的式（8-24）实现模型的损失函数。
# 判别器的loss为D(fake_samples)-D(real_samples)再加上联合分布样本的梯度惩罚项gradient_penalties，其中fake_samples为生成的模拟数据，real_Samples为真实数据，
# 生成器的loss为-D(fake_samples)。
def train(D,G,outdir,z_dimension,num_epochs=30):
    d_optimizer = torch.optim.Adam(D.parameters(),lr=0.001) # 定义优化器
    g_optimizer = torch.optim.Adam(G.parameters(),lr=0.001)
    os.makedirs(outdir,exist_ok=True) # 创建输出文件夹
    # 在函数train()中，判别器和生成器是分开训练的。让判别器学习的次数多一些，判别器每训练5次，生成器优化1次。
    # WGAN_gp不会因为判别器准确率太高而引起生成器梯度消失的问题，所以好的判别器会让生成器有更好的模拟效果。
    for epoch in range(num_epochs):
        for i,(img,lab) in enumerate(train_loader):
            num_img = img.size(0)
            # 训练判别器
            real_img = img.to(device)
            y_one_hot = torch.zeros(lab.shape[0],10).scatter_(1,lab.view(lab.shape[0],1),1).to(device)
            for ii in range(5): # 循环训练5次
                d_optimizer.zero_grad() # 梯度清零
                real_img = real_img.requires_grad_(True) # 在WGAN-gp基础上新增，将输入参数real_img设置为可导
                # 对real_img进行判别
                real_out = D(real_img, y_one_hot)
                # 生成随机值
                z = torch.randn(num_img, z_dimension).to(device)
                fake_img = G(z, y_one_hot)  # 生成fake_img
                fake_out = D(fake_img, y_one_hot)  # 对fake_img进行判别
                # 计算梯度惩罚项
                gradient_penalty_div = compute_w_div(real_img, real_out, fake_img, fake_out) # 使用gradient_penalty_div()求梯度
                # 计算判别器的loss
                d_loss = -torch.mean(real_out) + torch.mean(fake_out) + gradient_penalty_div
                d_loss.backward()
                d_optimizer.step()

            # 训练生成器
            for ii in range(1):
                g_optimizer.zero_grad()  # 梯度清0
                z = torch.randn(num_img, z_dimension).to(device)
                fake_img = G(z, y_one_hot)
                fake_out = D(fake_img, y_one_hot)
                g_loss = -torch.mean(fake_out)
                g_loss.backward()
                g_optimizer.step()
        # 输出可视化结果
        fake_images = to_img(fake_img.cpu().data)
        real_images = to_img(real_img.cpu().data)
        rel = torch.cat([to_img(real_images[:10]), fake_images[:10]], axis=0)
        imshow(torchvision.utils.make_grid(rel, nrow=10), os.path.join(outdir, 'fake_images-{}.png'.format(epoch + 1)))
        # 输出训练结果
        print('Epoch [{}/{}], d_loss: {:.6f}, g_loss: {:.6f} ''D real: {:.6f}, D fake: {:.6f}'.format(epoch, num_epochs,
                                                                                                      d_loss.data,
                                                                                                      g_loss.data,
                                                                                                      real_out.data.mean(),
                                                                                                      fake_out.data.mean()))
    # 保存训练模型
    torch.save(G.state_dict(), os.path.join(outdir, 'div-generator.pth'))
    torch.save(D.state_dict(), os.path.join(outdir, 'div-discriminator.pth'))

2.5コードコンバット：モデルの結果を視覚化----WGAN_div_241.py（パート5）

# 1.5 定义函数，实现可视化模型结果：获取一部分测试数据，显示由模型生成的模拟数据。
def displayAndTest(D,G,z_dimension):    # 可视化结果
    sample = iter(test_loader)
    images, labels = sample.next()
    y_one_hot = torch.zeros(labels.shape[0], 10).scatter_(1,labels.view(labels.shape[0], 1), 1).to(device)
    num_img = images.size(0) # 获取样本个数
    with torch.no_grad():
        z = torch.randn(num_img, z_dimension).to(device) # 生成随机数
        fake_img = G(z, y_one_hot)
    fake_images = to_img(fake_img.cpu().data) # 生成模拟样本
    rel = torch.cat([to_img(images[:10]), fake_images[:10]], axis=0)
    imshow(torchvision.utils.make_grid(rel, nrow=10))
    print(labels[:10])

2.6コードコンバット：関数の呼び出しとモデルのトレーニング----WGAN_div_241.py（パート6）

# 1.6 调用函数并训练模型：实例化判别器和生成器模型，并调用函数进行训练
if __name__ == '__main__':
    z_dimension = 40  # 设置输入随机数的维度

    D = WGAN_D().to(device)  # 实例化判别器
    G = WGAN_G(z_dimension).to(device)  # 实例化生成器
    train(D, G, './w_img', z_dimension) # 训练模型
    displayAndTest(D, G, z_dimension) # 输出可视化

結果：

WGAN-dIVモデルは、非常に明確なシミュレーションサンプルも出力します。WGAN-divに関する論文では、WGAN-divモデルとWGAN-gpモデルを比較したところ、WGAN-diIVモデルのFIDスコアが高いことがわかりました（FIDは生成された画像の品質を評価するための指標です） GANによる）。

 

3コードの概要（WGAN_div_241.py）

import torch
import torchvision
from torchvision import transforms
from torch.utils.data import DataLoader
from torch import nn
import torch.autograd as autograd
import matplotlib.pyplot as plt
import numpy as np
import matplotlib
import os
os.environ["KMP_DUPLICATE_LIB_OK"]="TRUE"

# 1.1 引入模块并载入样本:定义基本函数，加载FashionMNIST数据集
def to_img(x):
    x = 0.5 * (x+1)
    x = x.clamp(0,1)
    x = x.view(x.size(0),1,28,28)
    return x

def imshow(img,filename = None):
    npimg = img.numpy()
    plt.axis('off')
    array = np.transpose(npimg,(1,2,0))
    if filename != None:
        matplotlib.image.imsave(filename,array)
    else:
        plt.imshow(array)
        # plt.savefig(filename) # 保存图片 注释掉，因为会报错，暂时不知道什么原因 2022.3.26 15:20
        plt.show()

img_transform = transforms.Compose(
    [
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.5],std=[0.5])
    ]
)

data_dir = './fashion_mnist'

train_dataset = torchvision.datasets.FashionMNIST(data_dir,train=True,transform=img_transform,download=True)
train_loader = DataLoader(train_dataset,batch_size=1024,shuffle=True)
# 测试数据集
val_dataset = torchvision.datasets.FashionMNIST(data_dir,train=False,transform=img_transform)
test_loader = DataLoader(val_dataset,batch_size=10,shuffle=False)
# 指定设备
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
print(device)

# 1.2 实现生成器和判别器 ：因为复杂部分都放在loss值的计算方面了，所以生成器和判别器就会简单一些。
# 生成器和判别器各自有两个卷积和两个全连接层。生成器最终输出与输入图片相同维度的数据作为模拟样本。
# 判别器的输出不需要有激活函数，并且输出维度为1的数值用来表示结果。
# 在GAN模型中，因判别器的输入则是具体的样本数据，要区分每个数据的分布特征，所以判别器使用实例归一化，
class WGAN_D(nn.Module): # 定义判别器类D ：有两个卷积和两个全连接层
    def __init__(self,inputch=1):
        super(WGAN_D, self).__init__()
        self.conv1 = nn.Sequential(
            nn.Conv2d(inputch,64,4,2,1), # 输出形状为[batch,64,28,28]
            nn.LeakyReLU(0.2,True),
            nn.InstanceNorm2d(64,affine=True)
        )
        self.conv2 = nn.Sequential(
            nn.Conv2d(64,128,4,2,1),# 输出形状为[batch,64,14,14]
            nn.LeakyReLU(0.2,True),
            nn.InstanceNorm2d(128,affine=True)
        )
        self.fc = nn.Sequential(
            nn.Linear(128*7*7,1024),
            nn.LeakyReLU(0.2,True)
        )
        self.fc2 = nn.Sequential(
            nn.InstanceNorm1d(1,affine=True),
            nn.Flatten(),
            nn.Linear(1024,1)
        )
    def forward(self,x,*arg): # 正向传播
        x = self.conv1(x)
        x = self.conv2(x)
        x = x.view(x.size(0),-1)
        x = self.fc(x)
        x = x.reshape(x.size(0),1,-1)
        x = self.fc2(x)
        return x.view(-1,1).squeeze(1)

# 在GAN模型中，因生成器的初始输入是随机值，所以生成器使用批量归一化。
class WGAN_G(nn.Module): # 定义生成器类G：有两个卷积和两个全连接层
    def __init__(self,input_size,input_n=1):
        super(WGAN_G, self).__init__()
        self.fc1 = nn.Sequential(
            nn.Linear(input_size * input_n,1024),
            nn.ReLU(True),
            nn.BatchNorm1d(1024)
        )
        self.fc2 = nn.Sequential(
            nn.Linear(1024,7*7*128),
            nn.ReLU(True),
            nn.BatchNorm1d(7*7*128)
        )
        self.upsample1 = nn.Sequential(
            nn.ConvTranspose2d(128,64,4,2,padding=1,bias=False), # 输出形状为[batch,64,14,14]
            nn.ReLU(True),
            nn.BatchNorm2d(64)
        )
        self.upsample2 = nn.Sequential(
            nn.ConvTranspose2d(64,1,4,2,padding=1,bias=False), # 输出形状为[batch,64,28,28]
            nn.Tanh()
        )
    def forward(self,x,*arg): # 正向传播
        x = self.fc1(x)
        x = self.fc2(x)
        x = x.view(x.size(0),128,7,7)
        x = self.upsample1(x)
        img = self.upsample2(x)
        return img

# 1.3 计算w散度：返回值充当WGAN-gp中的惩罚项，用于计算判别器的损失
def compute_w_div(real_samples,real_out,fake_samples,fake_out):
    # 定义参数
    k = 2
    p = 6
    # 计算真实空间的梯度
    weight = torch.full((real_samples.size(0),),1,device=device)
    real_grad = autograd.grad(outputs=real_out,
                              inputs=real_samples,
                              grad_outputs=weight,
                              create_graph=True,
                              retain_graph=True,
                              only_inputs=True)[0]
    # L2范数
    real_grad_norm = real_grad.view(real_grad.size(0),-1).pow(2).sum(1)
    # 计算模拟空间的梯度
    fake_grad = autograd.grad(outputs=fake_out,
                              inputs=fake_samples,
                              grad_outputs=weight,
                              create_graph=True,
                              retain_graph=True,
                              only_inputs=True)[0]
    # L2范数
    fake_grad_norm = fake_grad.view(fake_grad.size(0),-1).pow(2).sum(1)
    # 计算W散度距离
    div_gp = torch.mean(real_grad_norm **(p/2)+fake_grad_norm**(p/2))*k/2
    return div_gp

## 1.4 定义模型的训练函数
# 定义函数train()，实现模型的训练过程。
# 在函数train()中，按照对抗神经网络专题（一）中的式（8-24）实现模型的损失函数。
# 判别器的loss为D(fake_samples)-D(real_samples)再加上联合分布样本的梯度惩罚项gradient_penalties，其中fake_samples为生成的模拟数据，real_Samples为真实数据，
# 生成器的loss为-D(fake_samples)。
def train(D,G,outdir,z_dimension,num_epochs=30):
    d_optimizer = torch.optim.Adam(D.parameters(),lr=0.001) # 定义优化器
    g_optimizer = torch.optim.Adam(G.parameters(),lr=0.001)
    os.makedirs(outdir,exist_ok=True) # 创建输出文件夹
    # 在函数train()中，判别器和生成器是分开训练的。让判别器学习的次数多一些，判别器每训练5次，生成器优化1次。
    # WGAN_gp不会因为判别器准确率太高而引起生成器梯度消失的问题，所以好的判别器会让生成器有更好的模拟效果。
    for epoch in range(num_epochs):
        for i,(img,lab) in enumerate(train_loader):
            num_img = img.size(0)
            # 训练判别器
            real_img = img.to(device)
            y_one_hot = torch.zeros(lab.shape[0],10).scatter_(1,lab.view(lab.shape[0],1),1).to(device)
            for ii in range(5): # 循环训练5次
                d_optimizer.zero_grad() # 梯度清零
                real_img = real_img.requires_grad_(True) # 在WGAN-gp基础上新增，将输入参数real_img设置为可导
                # 对real_img进行判别
                real_out = D(real_img, y_one_hot)
                # 生成随机值
                z = torch.randn(num_img, z_dimension).to(device)
                fake_img = G(z, y_one_hot)  # 生成fake_img
                fake_out = D(fake_img, y_one_hot)  # 对fake_img进行判别
                # 计算梯度惩罚项
                gradient_penalty_div = compute_w_div(real_img, real_out, fake_img, fake_out) # 求梯度
                # 计算判别器的loss
                d_loss = -torch.mean(real_out) + torch.mean(fake_out) + gradient_penalty_div
                d_loss.backward()
                d_optimizer.step()

            # 训练生成器
            for ii in range(1):
                g_optimizer.zero_grad()  # 梯度清0
                z = torch.randn(num_img, z_dimension).to(device)
                fake_img = G(z, y_one_hot)
                fake_out = D(fake_img, y_one_hot)
                g_loss = -torch.mean(fake_out)
                g_loss.backward()
                g_optimizer.step()
        # 输出可视化结果
        fake_images = to_img(fake_img.cpu().data)
        real_images = to_img(real_img.cpu().data)
        rel = torch.cat([to_img(real_images[:10]), fake_images[:10]], axis=0)
        imshow(torchvision.utils.make_grid(rel, nrow=10), os.path.join(outdir, 'fake_images-{}.png'.format(epoch + 1)))
        # 输出训练结果
        print('Epoch [{}/{}], d_loss: {:.6f}, g_loss: {:.6f} ''D real: {:.6f}, D fake: {:.6f}'.format(epoch, num_epochs,
                                                                                                      d_loss.data,
                                                                                                      g_loss.data,
                                                                                                      real_out.data.mean(),
                                                                                                      fake_out.data.mean()))
    # 保存训练模型
    torch.save(G.state_dict(), os.path.join(outdir, 'div-generator.pth'))
    torch.save(D.state_dict(), os.path.join(outdir, 'div-discriminator.pth'))

# 1.5 定义函数，实现可视化模型结果：获取一部分测试数据，显示由模型生成的模拟数据。
def displayAndTest(D,G,z_dimension):    # 可视化结果
    sample = iter(test_loader)
    images, labels = sample.next()
    y_one_hot = torch.zeros(labels.shape[0], 10).scatter_(1,labels.view(labels.shape[0], 1), 1).to(device)
    num_img = images.size(0) # 获取样本个数
    with torch.no_grad():
        z = torch.randn(num_img, z_dimension).to(device) # 生成随机数
        fake_img = G(z, y_one_hot)
    fake_images = to_img(fake_img.cpu().data) # 生成模拟样本
    rel = torch.cat([to_img(images[:10]), fake_images[:10]], axis=0)
    imshow(torchvision.utils.make_grid(rel, nrow=10))
    print(labels[:10])

# 1.6 调用函数并训练模型：实例化判别器和生成器模型，并调用函数进行训练
if __name__ == '__main__':
    z_dimension = 40  # 设置输入随机数的维度

    D = WGAN_D().to(device)  # 实例化判别器
    G = WGAN_G(z_dimension).to(device)  # 实例化生成器
    train(D, G, './w_img', z_dimension) # 训练模型
    displayAndTest(D, G, z_dimension) # 输出可视化