# 构建数据集迭代器
train_loader = torch.utils.data.DataLoader(train_dataset,
batch_size=batch_size,
shuffle=True)
test_loader = torch.utils.data.DataLoader(test_dataset,
batch_size=batch_size,
shuffle=True)
"""初始化CNN网络"""
Minionn_5 = Minionn_fivelayer(in_dim=1, n_class=10)
print('Minionn_5: \n', Minionn_5)
"""构建优化器"""
loss_fn = NLLLoss()
lr = 1e-3 # Adam优化器的学习率
beta1 = 0.5 # Adam优化器的参数（需要调整试试？）
optimizer = Adam(Minionn_5.parameters(),
learning_rate=lr,
betas=(beta1, 0.999)) # 测试一下Adam优化器
"""加载预训练的模型"""
pre_module_path = "./model_save/Minionn_5/Minionn_5_parameters-5.pkl"
params = torch.load(pre_module_path)
Minionn_5.load_state_dict(params['state_dict']) # 加载模型
n_epochs_pre = params['epoch']
"""迭代训练"""
for epoch in range(n_epochs):
# break
running_loss = 0.0
running_correct = 0
print("Epoch {}/{}".format(epoch, n_epochs))
print("-" * 10)
def test_dcgan():
start_time = time.time()
# 设置transform
cifar_transform = transforms.Compose([
Data_loader.cifar_Resize((3,32,32))
])
"""加载CIFAR-10数据集"""
root_dir = './data/cifar-10/cifar-10-python/cifar-10-batches-py'
cifar_train_dataset = Data_loader.CifarDataset(root_dir, transform=cifar_transform, train=True)
print('traindata_len: \n',len(cifar_train_dataset))
# 构建数据集迭代器
cifar_train_loader = torch.utils.data.DataLoader(cifar_train_dataset, batch_size=batch_size, shuffle=True)
"""初始化网络、参数"""
netG = Generator()
netG.apply(weights_init)
print('netG: \n', netG)
netD = Discriminator()
netD.apply(weights_init)
print('netD: \n', netD)
"""构建优化器"""
# 二进制交叉熵损失函数
loss = BECLoss()
# 噪声从标准正态分布（均值为0，方差为 1，即高斯白噪声）中随机抽取一组数
fixed_noise = np.random.normal(0.0, 1.2, size=(64,nz, 1,1)) # 用于G生成图像时固定的噪声初始化
# 定义真假样本标签
real_label = 1
fake_label = 0
# 定义Adam优化器
optimizerD = Adam(netD.parameters(), learning_rate=lr, betas=(beta1, 0.999))
optimizerG = Adam(netG.parameters(), learning_rate=lr, betas=(beta1, 0.999))
"""训练模型，生成数据"""
# 存储生成的图像
img_list = []
# 记录G和D的损失
G_losses = []
D_losses = []
iters = 0
"""加载预训练的模型"""
'''
pre_module_path = "./model_save/lenet_numpy_parameters-cifar-Adam-1.pkl"
params = torch.load(pre_module_path)
netD.load_state_dict(params['D_state_dict']) # 加载模型
netG.load_state_dict(params['G_state_dict']) # 加载模型
num_epochs_pre = params['epoch']
'''
print("----------start training loop----------")
for epoch in range(num_epochs):
# dataloader获取真实图像
for t, (data, target) in enumerate(cifar_train_loader, 0):
'''
(1)先更新D Update D network: minimize -[ log(D(x)) + log(1 - D(G(z))) ]
训练D的目标是让D更加有能力判断真假数据
'''
## 使用真实数据X进行训练（计算log(D(x))）
netD.zero_grad() # 训练更新前需要在每个batch中将梯度设置为0
real_data = data.detach().numpy()
# print('real_data: \n',real_data[0])
# print('real_data shape: \n',real_data[0].shape)
# break
## MNIST 数据需要先从1x28x28填充到1x32x32
# if is_mnist:
# real_data = np.pad(real_data, ((0, 0), (0, 0), (2, 2), (2, 2)), 'constant', constant_values=0)
b_size = real_data.shape[0]
label = np.full((b_size,), real_label)
# 计算D前向传播值
output_d_real = netD.forward(real_data).reshape(-1)
# 计算D真实数据交叉熵损失
errD_real = loss.forward(output_d_real, label)
# 计算D的梯度
dy_errD_real = loss.gradient()
netD.backward(dy_errD_real)
## 使用生成数据进行训练（计算log(1 - D(G(z)))）
noise = np.random.normal(0.0, 1.2, size=(b_size, nz, 1,1)) # 训练每次单独生成噪声
# G生成假数据
fake_data = netG.forward(noise)
label.fill(fake_label)
# D识别假数据
output_d_fake = netD.forward(fake_data).reshape(-1)
# 计算D假数据交叉熵损失
errD_fake = loss.forward(output_d_fake, label)
# 计算D的梯度
dy_errD_fake = loss.gradient()
netD.backward(dy_errD_fake)
# 计算总损失
errD = errD_real+errD_fake
# 计算D(x),D(G(z))的均值
D_x = np.mean(output_d_real)
D_G_z1 = np.mean(output_d_fake)
# 更新D参数
optimizerD.step()
'''
(2)更新G Update G network: minimize -log(D(G(z)))
'''
netG.zero_grad()
# 填充真实标签，使得交叉熵函数可以只计算log(D(G(z))部分
label.fill(real_label)
output_d_fake = netD.forward(fake_data).reshape(-1)
errG = loss.forward(output_d_fake, label)
# 计算G的梯度（梯度需要从D传向G）
dy_errG = loss.gradient()
dy_netD = netD.backward(dy_errG)
netG.backward(dy_netD)
# 计算D(G(z))的均值
D_G_z2 = np.mean(output_d_fake)
# 更新G参数（不会去计算D的梯度2333）
optimizerG.step()
"""输出训练状态"""
# Loss_D
# Loss_G
# D(x)：训练中D对真实数据的平均预测输出
# D(G(z))：训练中D对虚假数据的平均预测输出（为啥是除法？？）
if t % 10 == 0:
end_time1 = time.time()
print('[%d/%d][%d/%d]\t Loss_D: %.4f\t Loss_G: %.4f\t D(x): %.4f\t D(G(z)): %.4f / %.4f\t train time: %.4f min'
% (epoch, num_epochs, t, len(cifar_train_loader), errD, errG, D_x, D_G_z1, D_G_z2, (end_time1-start_time)/60))
# 记录损失的历史，可以用作画图
G_losses.append(errG)
D_losses.append(errD)
# 记录G生成的图像
if (iters % 500 == 0) or ((epoch == num_epochs-1) and (t == len(cifar_train_loader)-1)):
fake_img = netG.forward(fixed_noise)# 一次生成64张图
fake_tensor = torch.tensor(fake_img)
img_list.append(vutils.make_grid(fake_tensor, padding=2, normalize=True))
iters += 1
"""绘图：记录损失"""
plt.figure(figsize=(10,5))
plt.title("Generator and Discriminator Loss During Training")
plt.plot(G_losses,label="G")
plt.plot(D_losses,label="D")
plt.xlabel("iterations")
plt.ylabel("Loss")
plt.legend()
# 保存图片
time_stemp = time.strftime("%Y-%m-%d-%H-%M", time.localtime())
plt.savefig('./experiment_img/gan_generate/Loss_fig-cifar-Adam-'+str(epoch+num_epochs_pre)+'('+time_stemp+').png')
# plt.show()
"""绘图：记录G输出"""
real_batch = next(iter(cifar_train_loader))
# Plot the real images
plt.figure(figsize=(15,15))
plt.subplot(1,2,1)
plt.axis("off")
plt.title("Real Images")
plt.imshow(np.transpose(vutils.make_grid(real_batch[0][:64], padding=5, normalize=True).cpu(),(1,2,0)))
# Plot the fake images from the last epoch
plt.subplot(1,2,2)
plt.axis("off")
plt.title("Fake Images")
plt.imshow(np.transpose(img_list[-1],(1,2,0)))
# 保存图片
plt.savefig('./experiment_img/gan_generate/Real_Generate-cifar-Adam-'+str(epoch+num_epochs_pre)+'('+time_stemp+').png')
# plt.show()
end_time = time.time()
print('training time: \n', (end_time-start_time)/60)
'''存储模型'''
checkpoint_path = "./model_save/DCGAN_numpy_parameters-cifar-Adam-"+str(epoch+num_epochs_pre)+".pkl"
torch.save({'epoch':num_epochs+num_epochs_pre, 'D_state_dict':netD.state_dict(), 'G_state_dict':netG.state_dict(), 'G_losses':G_losses, 'D_losses':D_losses}, checkpoint_path)