from torch.utils.data import DataLoader
from torchvision.datasets import CIFAR10
import torchvision.transforms as transforms
from torchvision.utils import save_image

dataset=CIFAR10(root='./data', download=False, transform=transforms.ToTensor())
dataloader=DataLoader(dataset,batch_size=64,shuffle=True)

for batch_idx,data in enumerate(dataloader):
    real_images,_=data
    batch_size=real_images.size()
    print('#{} has {} images :'.format(batch_idx,batch_size))
    if batch_idx%100==0:
        path='./data/CIFAR10_shuffled_batch{:03d}.png'.format(batch_idx)
        save_image(real_images,path,normalize=True)

import torch.nn as nn

#搭建生成网络
latent_size=64 #潜在大小
n_channel=3    #输出通道数
n_g_feature=64 #生成网络隐藏层大小
gnet=nn.Sequential(
     #输出大小（64,1,1）
     nn.ConvTranspose2d(latent_size,4*n_g_feature,kernel_size=4,bias=False),
     nn.BatchNorm2d(4*n_g_feature),
     nn.ReLU(),
     #输出大小（256,4,4）
     nn.ConvTranspose2d(4*n_g_feature,2*n_g_feature,kernel_size=4,stride=2,padding=1,bias=False),
     nn.BatchNorm2d(2*n_g_feature),
     nn.ReLU(),
     #输出大小（128，8，8）
     nn.ConvTranspose2d(2*n_g_feature,n_g_feature,kernel_size=4,stride=2,padding=1,bias=False),
     nn.BatchNorm2d(n_g_feature),
     nn.ReLU(),
     #输出大小(64,16,16)
     nn.ConvTranspose2d(n_g_feature,n_channel,kernel_size=4,stride=2,padding=1),
     nn.Sigmoid())
    #图片大小(3,32,32)
print(gnet)
#搭建鉴别网络
n_d_feature = 64#鉴别网络隐藏层大小
dnet=nn.Sequential(
    #图片大小(3,32,32)
     nn.Conv2d(n_channel,n_d_feature,kernel_size=4,stride=2,padding=1),
     nn.LeakyReLU(0.2),
    #大小(64,16,16)
     nn.Conv2d(n_d_feature,2*n_d_feature,kernel_size=4,stride=2,padding=1,bias=False),
     nn.BatchNorm2d(2*n_d_feature),
     nn.LeakyReLU(0.2),
    #大小(128,8,8)
     nn.Conv2d(2*n_d_feature,4*n_d_feature,kernel_size=4,stride=2,padding=1,bias=False),
     nn.BatchNorm2d(4*n_d_feature),
     nn.LeakyReLU(0.2),
     #大小（256，4,4）
     nn.Conv2d(4*n_d_feature,1,kernel_size=4),
    #对数赔率张量大小(1,1,1))
)
print(dnet)

import torch.nn.init as init

#初始化权重值的函数
def weights_init(m):
    if type(m) in [nn.ConvTranspose2d,nn.Conv2d]:
        init.xavier_normal_(m.weight)
    elif type(m)==nn.BatchNorm2d:
        init.normal_(m.weight,1.0,0.02)
        init.constant_(m.bias,0)
gnet.apply(weights_init)
dnet.apply(weights_init)
    

import torch
import torch.optim

#损失
criterion=nn.BCEWithLogitsLoss();
 
#优化器
goptimizer=torch.optim.Adam(gnet.parameters(),lr=0.002,betas=(0.5,0.999))
doptimizer=torch.optim.Adam(dnet.parameters(),lr=0.002,betas=(0.5,0.999))

#用于测试的固定噪声，用来查看相同的潜在张量在训练过程中生成图片的变换
batch_size=64
fixed_noises=torch.randn(batch_size,latent_size,1,1)

#训练过程
epoch_num=10
for epoch in range(epoch_num):
    for batch_idx,data in enumerate(dataloader):
        # 载入本批次数据
        real_images,_=data
        batch_size=real_images.size(0)
        
        # 训练鉴别网络
        labels=torch.ones(batch_size)#真实数据标签为1
        preds=dnet(real_images)#对真实数据进行判别
        outputs=preds.reshape(-1)
        dloss_real=criterion(outputs,labels)#真实数据对应标签为1
        dmean_real=outputs.sigmoid().mean()
        #计算鉴别器将多少比例的真实数据判定为真，仅用于输出显示
        
        noises=torch.randn(batch_size,latent_size,1,1)#潜在噪声
        fake_images=gnet(noises)#生成假数据
        labels=torch.zeros(batch_size)#假数据对应标签为0
        fake=fake_images.detach()
        #使得梯度计算不回溯到生成网络，可用于加快训练速度。删去此步骤，结果不变
        
        preds=dnet(fake)#对假数据进行鉴别
        outputs=preds.view(-1)
        dloss_fake=criterion(outputs,labels)#假数据的鉴别损失
        
        dmean_fake=outputs.sigmoid().mean()
        #计算鉴别器将多少比例的假数据判别为真 仅用于输出显示
        
        dloss=dloss_real+dloss_fake #总鉴别器的损失
        dnet.zero_grad()
        dloss.backward()
        doptimizer.step()
        
        #训练生成网络
        labels=torch.ones(batch_size)
        #生成网络希望所有数据都被认为是真数据
        preds=dnet(fake_images)#让假数据通过鉴别网络
        outputs=preds.view(-1)
        gloss=criterion(outputs,labels)#从真数据中看到损失
        gmean_fake=outputs.sigmoid().mean()
        #计算鉴别器将多少比例的数据判定为真，仅用于输出显示
        gnet.zero_grad()
        gloss.backward()
        goptimizer.step()
        
        #输出本步训练结果
        print('[{}/{}]'.format(epoch,epoch_num)+'[{}/{}]'.format(batch_idx,len(dataloader))+
             '鉴别网络损失：{:g} 生成网络损失：{:g}'.format(dloss,gloss)+
             '真实数据判真比例：{:g} 假数据判真比例：{:g}/{:g}'.format(dmean_real,dmean_fake,gmean_fake))
        if batch_idx%100==0:
            fake=gnet(fixed_noises)#由固定潜在张量生成数据
            save_image(fake,'./data/image_epoch{:02d}_batch{:03d}.png'.format(epoch,batch_idx))
            #保存假数据
        
        
        

import jovian
jovian.commit()