1 WGAN-gp模型生成模拟数据案例说明
使用WGAN-gp模型模拟Fashion-MNIST数据的生成，会使用到WGAN-gp模型、深度卷积GAN(DeepConvolutional GAN，DCGAN)模型、实例归一化技术。
1.1 DCGAN中的全卷积
WGAN-gp模型侧重于GAN模型的训练部分，而DCGAN是指使用卷积神经网络的GAN，它侧重于GAN模型的结构部分，重点介绍在DCGAN中使用全卷积进行重构的技术。
1.1.1 DCGAN的原理与实现
DCGAN的原理和GAN类似，只是把CNN卷积技术用在GAN模式的网络里。生成器G在生成数据时，使用反卷积的重构技术来重构原始图片，判别器D使用卷积技术来识别图片特征，进而做出判别。同时，DCGAN中的卷积神经网络对结构进行改变，提高样本质量与收敛速度。
- G网络中取消所有池化层，使用全卷积，并且采用大于或等于2的步长来进行上采样，使用ReLU()作为激活函数，最后一层使用tanh()。
- D网络中用加入下采样的卷积操作代替池化，通常会使用LeakyReLU()作为激活函数，这种激活函数可以对小于0的部分特征给予保留。
- 通常不在D和G的最后一层使用归一化处理，这样做的目的是保证模型能够学习到数据的正确分布。
DCGAN模型可以更好地学到对输入图像层次化的表示，尤其在生成器部分会有更好的模拟效果，在训练中，会使用Adam优化算法。
1.1.2 全卷积的实现
在PyTorch中，全卷积是通过转置卷积接口ConvTranspose2d()来实现的。该接口的参数与卷积函数参数的含义相同，还有1D和3D的转置卷积实现。
convTranspose2d (in_channels，out_channels，kernel_size，stride=1，padding=0,output_padding=0，groups=1，bias=τrue，dilation=1，padding_mode='zeros')先对卷积核进行转置，再实现全卷积处理，输出的尺寸与卷积操作所输出的尺寸互逆。
1.2 上采样与下采样
1.2.1 上下采样的简述
- 上采样是将图像放大。
- 下采样是将图像缩小。
上采样与下采样操作并不能给图片带来更多的信息，但会对图像质量产生影响。
在深度卷积网络模型的运算中，通过上采样与下采样操作可实现本层数据与上下层的维度匹配。
1.2.2 上下采样的作用
神经网络模型常使用窄卷积或池化对模型进行下采样，
神经网络模型使用转置卷积对模型进行上采样。
1.2.3 通过卷积和全卷积函数实现下采样方式处理再上采样的方式进行还原操作。
from torch import nn
import torch
# 定义输入数据,3通道,尺寸为[12,12]
input = torch.randn(1,3,12,12)
# 输入和输出通道为3,卷积核为3,步长为2,进行下采样
downsample = nn.Conv2d(3,3,3,stride=2,padding=1)
h = downsample(input)
print(h.size())
# 输出结果:torch.SizeC[1,3,6,6]),尺寸变为[6,6]
# 输入和输出通道为3,卷积核为3,步长为2,进行上采样还原
upsample = nn.ConvTranspose2d(3,3,3,stride=2,padding=1)
output = upsample(h,output_size=input.size())
print(output.size())
# 输出结果:torch.Size([1,3,12,12]),尺寸变回[12,12]1.3 实例归一化
批量归一化是对一个批次图片中的所有像素求均值和标准差，
实例归一化是对单一图片进行归一化处理，即对单个图片的所有像素求均值和标准差
1.3.1 实例归一化的使用场景
在对抗神经网络模型、风格转换这类生成式任务中，常用实例归一化取代批量归一化，因为生成式任务的本质是将生成样本的特征分布与目标样本的特征分布进行匹配。生成式任务的每个样本都有独立的风格，不应该与批次中其他样本产生太多联系。因此，实例归一化适合解决这种基于个体的样本分布问题。
1.3.2 实例归一化的使用
PyTorch中实例归一化的实现接口是nn模块下的InstanceNorm2d()，还有1D、3D实例归一化，与该接口类似。
Instanceormd(num_features,eps=1e-5,momentum=0.1,affine=False,track_running_stats=False)仅关注参数num_features，该参数是需要传入输入数据的通道数，其他参数与批量归一化BatchNorm2d()中参数的含义一致，该接口会按照通道对单个数据进行归一化，其返回的形状与输入形状相同。
2 实例代码编写
2.1 代码实战：引入模块并载入样本----WGAN-gp-228.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-gp-228.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 img2.3 代码实战：定义函数完成梯度惩罚项----WGAN-gp-228.py（第3部分）
# 1.3 定义函数compute_gradient_penalty()完成梯度惩罚项
# 惩罚项的样本X_inter由一部分Pg分布和一部分Pr分布组成，同时对D(X_inter)求梯度，并计算梯度与1的平方差，最终得到gradient_penalties
lambda_gp = 10
# 计算梯度惩罚项
def compute_gradient_penalty(D,real_samples,fake_samples,y_one_hot):
# 获取一个随机数，作为真假样本的采样比例
eps = torch.FloatTensor(real_samples.size(0),1,1,1).uniform_(0,1).to(device)
# 按照eps比例生成真假样本采样值X_inter
X_inter = (eps * real_samples + ((1-eps)*fake_samples)).requires_grad_(True)
d_interpolates = D(X_inter,y_one_hot)
fake = torch.full((real_samples.size(0),),1,device=device) # 计算梯度输出的掩码，在本例中需要对所有梯度进行计算，故需要按照样本个数生成全为1的张量。
# 求梯度
gradients = autograd.grad(outputs=d_interpolates, # 输出值outputs，传入计算过的张量结果
inputs=X_inter,# 待求梯度的输入值inputs，传入可导的张量，即requires_grad=True
grad_outputs=fake, # 传出梯度的掩码grad_outputs，使用1和0组成的掩码，在计算梯度之后，会将求导结果与该掩码进行相乘得到最终结果。
create_graph=True,
retain_graph=True,
only_inputs=True
)[0]
gradients = gradients.view(gradients.size(0),-1)
gradient_penaltys = ((gradients.norm(2, dim=1) - 1) ** 2).mean() * lambda_gp
return gradient_penaltys2.4 代码实战：定义模型的训练函数----WGAN-gp-228.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_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 = compute_gradient_penalty(D,real_img.data,fake_img.data,y_one_hot)
# 计算判别器的loss
d_loss = -torch.mean(real_out)+torch.mean(fake_out)+gradient_penalty
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, 'generator.pth'))
torch.save(D.state_dict(), os.path.join(outdir, 'discriminator.pth'))2.5 代码实战：现可视化模型结果----WGAN-gp-228.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-gp-228.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) # 输出可视化可以看到g_loss的绝对值在逐渐变小，d_loss的绝对值在逐渐变大。这表明生成的横拟样本质量越来越高。在本地路径的w_img文件夹下，可以看到30张图片，这里列出3张。
显示训练过程中的3张图片（每两行为一张)，它们分别是第1次、第18次和第30次迭代训练后的输出结果。每张图片的第1行为样本数据，第2行为生成的模拟数据。
得出结论：在WGAN-gp的判别器严格要求下，生成器生成的模拟数据越来越逼真。从生成的结果中可以看出，样本数据与生成的模拟数据类别并不对应，这是因为我们没有对其加入生成类别的信息。使用条件GAN可实现类别对应的效果。
3 代码汇总（WGAN-gp-228.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 定义函数compute_gradient_penalty()完成梯度惩罚项
# 惩罚项的样本X_inter由一部分Pg分布和一部分Pr分布组成，同时对D(X_inter)求梯度，并计算梯度与1的平方差，最终得到gradient_penalties
lambda_gp = 10
# 计算梯度惩罚项
def compute_gradient_penalty(D,real_samples,fake_samples,y_one_hot):
# 获取一个随机数，作为真假样本的采样比例
eps = torch.FloatTensor(real_samples.size(0),1,1,1).uniform_(0,1).to(device)
# 按照eps比例生成真假样本采样值X_inter
X_inter = (eps * real_samples + ((1-eps)*fake_samples)).requires_grad_(True)
d_interpolates = D(X_inter,y_one_hot)
fake = torch.full((real_samples.size(0),),1,device=device) # 计算梯度输出的掩码，在本例中需要对所有梯度进行计算，故需要按照样本个数生成全为1的张量。
# 求梯度
gradients = autograd.grad(outputs=d_interpolates, # 输出值outputs，传入计算过的张量结果
inputs=X_inter,# 待求梯度的输入值inputs，传入可导的张量，即requires_grad=True
grad_outputs=fake, # 传出梯度的掩码grad_outputs，使用1和0组成的掩码，在计算梯度之后，会将求导结果与该掩码进行相乘得到最终结果。
create_graph=True,
retain_graph=True,
only_inputs=True
)[0]
gradients = gradients.view(gradients.size(0),-1)
gradient_penaltys = ((gradients.norm(2, dim=1) - 1) ** 2).mean() * lambda_gp
return gradient_penaltys
# 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_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 = compute_gradient_penalty(D,real_img.data,fake_img.data,y_one_hot)
# 计算判别器的loss
d_loss = -torch.mean(real_out)+torch.mean(fake_out)+gradient_penalty
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, 'generator.pth'))
torch.save(D.state_dict(), os.path.join(outdir, '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) # 输出可视化