生成网络概括

VAE

# coding=utf-8
import torch.autograd
import torch.nn as nn
from torchvision import transforms
from torchvision import datasets
from torchvision.utils import save_image
import torch.nn.functional as F
import matplotlib.pyplot as plt
import os
from torchvision.utils import make_grid
# 创建文件夹
if not os.path.exists('./img_VAE'):
    os.mkdir('./img_VAE')
# GPU
device = 'cuda' if torch.cuda.is_available() else 'cpu'

def to_img(x):
    # out = 0.5 * (x+0.5)
    img = make_grid(x, nrow=8, normalize=True).detach()
    # out = x.clamp(0, 1)  # Clamp函数可以将随机变化的数值限制在一个给定的区间[min, max]内:
    # out = out.view(-1, 1, 28, 28)  # view()函数作用是将一个多行的Tensor,拼接成一行
    return img

batch_size = 64
num_epoch = 15
z_dimension = 2
# 图形啊处理过程
img_transform = transforms.Compose([
    transforms.ToTensor(),
])

# mnist dataset mnist数据集下载
mnist = datasets.MNIST(root='./data/', train=True, transform=img_transform, download=True)

# data loader 数据载入
dataloader = torch.utils.data.DataLoader(dataset=mnist, batch_size=batch_size, shuffle=True)

# 定义判别器  #####Discriminator######使用多层网络来作为判别器
class VAE(nn.Module):
    def __init__(self):
        super(VAE, self).__init__()
        # 定义编码器
        self.encoder = nn.Sequential(
            nn.Conv2d(1,16,kernel_size=3,stride=2,padding=1),
            nn.BatchNorm2d(16),
            nn.LeakyReLU(0.2,inplace=True),
            nn.Conv2d(16,32,kernel_size=3,stride=2,padding=1),
            nn.BatchNorm2d(32),
            nn.LeakyReLU(0.2,inplace=True),
            nn.Conv2d(32,32,kernel_size=3,stride=1,padding=1),
            nn.BatchNorm2d(32),
            nn.LeakyReLU(0.2,inplace=True),
        )
        self.encoder_fc1=nn.Linear(32*7*7,z_dimension)
        self.encoder_fc2=nn.Linear(32*7*7,z_dimension)
        self.Sigmoid = nn.Sigmoid()
        self.decoder_fc = nn.Linear(z_dimension,32 * 7 * 7)
        self.decoder = nn.Sequential(
            nn.ConvTranspose2d(32, 16, 4, 2, 1),
            nn.ReLU(inplace=True),
            nn.ConvTranspose2d(16, 1, 4, 2, 1),
            nn.Sigmoid(),
        )

    def noise_reparameterize(self,mean,logvar):
        eps = torch.randn(mean.shape).to(device)
        z = mean + eps * torch.exp(logvar)
        return z

    def forward(self, x):
        out1,out2 = self.encoder(x),self.encoder(x)
        mean = self.encoder_fc1(out1.view(out1.shape[0],-1))
        logstd = self.encoder_fc2(out2.view(out2.shape[0],-1))
        z = self.noise_reparameterize(mean,logstd)
        out3 = self.decoder_fc(z)
        out3 = out3.view(out3.shape[0],32,7,7)
        out3 = self.decoder(out3)
        return out3,mean,logstd

def loss_function(recon_x,x,mean,std):
    BCE = F.binary_cross_entropy(recon_x,x,reduction='sum')
    # 因为var是标准差的自然对数,先求自然对数然后平方转换成方差
    var = torch.pow(torch.exp(std),2)
    KLD = -0.5 * torch.sum(1+torch.log(var)-torch.pow(mean,2)-var)
    return BCE+KLD
# 创建对象
vae = VAE().to(device)
# vae.load_state_dict(torch.load('./VAE_z2.pth'))
vae_optimizer = torch.optim.Adam(vae.parameters(), lr=0.0003,
                                 betas=(0.9, 0.999), eps=1e-08, weight_decay=0)
###########################进入训练##判别器的判断过程#####################
for epoch in range(num_epoch):  # 进行多个epoch的训练
    for i, (img, _) in enumerate(dataloader):
        num_img = img.size(0)
        # view()函数作用把img变成[batch_size,channel_size,784]
        img = img.view(num_img,  1,28,28).to(device)  # 将图片展开为28*28=784
        x,mean,std = vae(img)  # 将真实图片放入判别器中
        loss = loss_function(x,img,mean,std)
        vae_optimizer.zero_grad()  # 在反向传播之前,先将梯度归0
        loss.backward()  # 将误差反向传播
        vae_optimizer.step()  # 更新参数
        # try:
        if (i + 1) % 100 == 0:
            print('Epoch[{}/{}],vae_loss:{:.6f} '.format(
                epoch, num_epoch, loss.item(),
            ))

        if epoch == 0:
            real_images = make_grid(img.cpu(), nrow=8, normalize=True).detach()
            save_image(real_images, './img_VAE/real_images.png')
        sample = torch.randn(64,z_dimension).to(device)
        output = vae.decoder_fc(sample)
        output = vae.decoder(output.view(output.shape[0],32,7,7))
        fake_images = make_grid(x.cpu(), nrow=8, normalize=True).detach()
        save_image(fake_images, './img_VAE/fake_images-{}.png'.format(epoch + 16))
# 保存模型
torch.save(vae.state_dict(), './VAE_z2.pth')

GAN

# coding=utf-8
import torch.autograd
import torch.nn as nn
from torchvision import transforms
from torchvision import datasets
from torchvision.utils import save_image
import os

# 创建文件夹
if not os.path.exists('./img_GAN'):
    os.mkdir('./img_GAN')
# GPU
device = 'cuda' if torch.cuda.is_available() else 'cpu'
def to_img(x):
    out = 0.5 * (x + 1)
    out = out.clamp(0, 1)  # Clamp函数可以将随机变化的数值限制在一个给定的区间[min, max]内:
    out = out.view(-1, 1, 28, 28)  # view()函数作用是将一个多行的Tensor,拼接成一行
    return out

batch_size = 128
num_epoch = 100
z_dimension = 100

# 图形啊处理过程
img_transform = transforms.Compose([
    transforms.ToTensor(),
    # transforms.Lambda(lambda x: x.repeat(3,1,1)),
    transforms.Normalize(mean=[0.5], std=[0.5])
])

# mnist dataset mnist数据集下载
mnist = datasets.MNIST(
    root='./data/', train=True, transform=img_transform, download=True
)

# data loader 数据载入
dataloader = torch.utils.data.DataLoader(
    dataset=mnist, batch_size=batch_size, shuffle=True
)

# 定义判别器  #####Discriminator######使用多层网络来作为判别器

# 将图片28x28展开成784,然后通过多层感知器,中间经过斜率设置为0.2的LeakyReLU激活函数,
# 最后接sigmoid激活函数得到一个0到1之间的概率进行二分类。
class discriminator(nn.Module):
    def __init__(self):
        super(discriminator, self).__init__()
        self.dis = nn.Sequential(
            nn.Linear(784, 256),  # 输入特征数为784,输出为256
            nn.LeakyReLU(0.2),  # 进行非线性映射
            nn.Linear(256, 256),  # 进行一个线性映射
            nn.LeakyReLU(0.2),
            nn.Linear(256, 1),
            nn.Sigmoid()  # 也是一个激活函数,二分类问题中,
            # sigmoid可以班实数映射到【0,1】,作为概率值,
            # 多分类用softmax函数
        )

    def forward(self, x):
        x = self.dis(x)
        return x

####### 定义生成器 Generator #####
# 输入一个100维的0~1之间的高斯分布,然后通过第一层线性变换将其映射到256维,
# 然后通过LeakyReLU激活函数,接着进行一个线性变换,再经过一个LeakyReLU激活函数,
# 然后经过线性变换将其变成784维,最后经过Tanh激活函数是希望生成的假的图片数据分布
# 能够在-1~1之间。
class generator(nn.Module):
    def __init__(self):
        super(generator, self).__init__()
        self.gen = nn.Sequential(
            nn.Linear(100, 256),  # 用线性变换将输入映射到256维
            nn.ReLU(True),  # relu激活
            nn.Linear(256, 256),  # 线性变换
            nn.ReLU(True),  # relu激活
            nn.Linear(256, 784),  # 线性变换
            nn.Tanh()  # Tanh激活使得生成数据分布在【-1,1】之间
        )

    def forward(self, x):
        x = self.gen(x)
        return x

# 创建对象
D = discriminator()
G = generator()
D = D.to(device)
G = G.to(device)

#########判别器训练train#####################
# 分为两部分:1、真的图像判别为真;2、假的图像判别为假
# 此过程中,生成器参数不断更新

# 首先需要定义loss的度量方式  (二分类的交叉熵)
# 其次定义 优化函数,优化函数的学习率为0.0003
criterion = nn.BCELoss()  # 是单目标二分类交叉熵函数
d_optimizer = torch.optim.Adam(D.parameters(), lr=0.0003)
g_optimizer = torch.optim.Adam(G.parameters(), lr=0.0003)

###########################进入训练##判别器的判断过程#####################

for epoch in range(num_epoch):  # 进行多个epoch的训练
    for i, (img, _) in enumerate(dataloader):
        for a in range(10):
            num_img = img.size(0)
            # view()函数作用把img变成[batch_size,channel_size,784]
            img = img.view(num_img,  -1)  # 将图片展开为28*28=784
            real_img = img.to(device) # 将tensor变成Variable放入计算图中
            real_label = torch.ones(num_img).to(device)# 定义真实的图片label为1
            fake_label = torch.zeros(num_img).to(device)  # 定义假的图片的label为0
            # print(img.shape)
            # 计算真实图片的损失
            real_out = D(real_img)  # 将真实图片放入判别器中
            d_loss_real = criterion(real_out, real_label)  # 得到真实图片的loss
            real_scores = real_out  # 得到真实图片的判别值,输出的值越接近1越好

            # 计算假的图片的损失
            z = torch.randn(num_img, z_dimension).to(device)  # 随机生成一些噪声
            fake_img = G(z)  # 随机噪声放入生成网络中,生成一张假的图片
            fake_out = D(fake_img)  # 判别器判断假的图片
            d_loss_fake = criterion(fake_out, fake_label)  # 得到假的图片的loss
            fake_scores = fake_out  # 得到假图片的判别值,对于判别器来说,假图片的损失越接近0越好

            # 损失函数和优化
            d_loss = d_loss_real + d_loss_fake  # 损失包括判真损失和判假损失
            d_optimizer.zero_grad()  # 在反向传播之前,先将梯度归0
            d_loss.backward()  # 将误差反向传播
            d_optimizer.step()  # 更新参数

        # ==================训练生成器============================
        ################################生成网络的训练###############################
        # 原理:目的是希望生成的假的图片被判别器判断为真的图片,
        # 在此过程中,将判别器固定,将假的图片传入判别器的结果与真实的label对应,
        # 反向传播更新的参数是生成网络里面的参数,
        # 这样可以通过更新生成网络里面的参数,来训练网络,使得生成的图片让判别器以为是真的
        # 这样就达到了对抗的目的

        # 计算假的图片的损失

        z = torch.randn(num_img, z_dimension).to(device) # 得到随机噪声
        fake_img = G(z)  # 随机噪声输入到生成器中,得到一副假的图片
        output = D(fake_img).squeeze(1)  # 经过判别器得到的结果
        g_loss = criterion(output, real_label)  # 得到的假的图片与真实的图片的label的loss

        # bp and optimize
        g_optimizer.zero_grad()  # 梯度归0
        g_loss.backward()  # 进行反向传播
        g_optimizer.step()  # .step()一般用在反向传播后面,用于更新生成网络的参数

        # 打印中间的损失
        # try:
        if (i + 1) % 100 == 0:
            print('Epoch[{}/{}],d_loss:{:.6f},g_loss:{:.6f} '
                  'D real: {:.6f},D fake: {:.6f}'.format(
                epoch, num_epoch, d_loss.item(), g_loss.item(),
                torch.mean(real_scores).item(), torch.mean(fake_scores).item()  # 打印的是真实图片的损失均值
            ))
        # except BaseException as e:
        #     pass

        if epoch == 0:
            real_images = to_img(real_img.cpu().data)
            save_image(real_images, './img_GAN/real_images.png')

        fake_images = to_img(fake_img.cpu().data)
        save_image(fake_images, './img_GAN/fake_images-{}.png'.format(epoch + 1))
# 保存模型
torch.save(G.state_dict(), './generator_GAN.pth')
torch.save(D.state_dict(), './discriminator_GAN.pth')

CGAN

# coding=utf-8
import torch.autograd
import torch.nn as nn
from torchvision import transforms
from torchvision import datasets
from torchvision.utils import save_image
import os

# 创建文件夹
if not os.path.exists('./img_CGAN'):
    os.mkdir('./img_CGAN')
# GPU
device = 'cuda' if torch.cuda.is_available() else 'cpu'
batch_size = 128
num_epoch = 25
z_dimension = 100
# 图形啊处理过程
img_transform = transforms.Compose([
    transforms.ToTensor(),
])
# mnist dataset mnist数据集下载
mnist = datasets.MNIST(
    root='./data/', train=True, transform=img_transform, download=True
)
# data loader 数据载入
dataloader = torch.utils.data.DataLoader(
    dataset=mnist, batch_size=batch_size, shuffle=True
)
# 定义判别器  #####Discriminator######使用多层网络来作为判别器
# 将图片28x28展开成784,然后通过多层感知器,中间经过斜率设置为0.2的LeakyReLU激活函数,
class discriminator(nn.Module):
    def __init__(self):
        super(discriminator, self).__init__()
        self.dis = nn.Sequential(
            nn.Linear(784, 256),  # 输入特征数为784,输出为256
            nn.LeakyReLU(0.2),  # 进行非线性映射
            nn.Linear(256, 256),  # 进行一个线性映射
            nn.LeakyReLU(0.2),
            nn.Linear(256, 10),
            nn.Softmax()
        )
    def forward(self, x):
        x = self.dis(x)
        return x
####### 定义生成器 Generator #####
class generator(nn.Module):
    def __init__(self):
        super(generator, self).__init__()
        self.gen = nn.Sequential(
            nn.Linear(z_dimension+10, 256),  # 用线性变换将输入映射到256维
            nn.ReLU(True),  # relu激活
            nn.Linear(256, 256),  # 线性变换
            nn.ReLU(True),  # relu激活
            nn.Linear(256, 784),  # 线性变换
            nn.Tanh()  # Tanh激活使得生成数据分布在【-1,1】之间
        )
    def forward(self, x):
        x = self.gen(x)
        return x
# 创建对象
D = discriminator()
G = generator()
D = D.to(device)
G = G.to(device)
# 载入模型
# G.load_state_dict(torch.load('./generator_CGAN_z100.pth'))
# D.load_state_dict(torch.load('./discriminator_CGAN_z100.pth'))
#########判别器训练train#####################
criterion = nn.BCELoss()  # 是单目标二分类交叉熵函数
d_optimizer = torch.optim.Adam(D.parameters(), lr=0.0003)
g_optimizer = torch.optim.Adam(G.parameters(), lr=0.0003)
###########################进入训练##判别器的判断过程#####################
for epoch in range(num_epoch):  # 进行多个epoch的训练
    for i, (img, label) in enumerate(dataloader):
        num_img = img.size(0)
        label_onehot = torch.zeros((num_img,10)).to(device)
        label_onehot[torch.arange(num_img),label]=1
        # view()函数作用把img变成[batch_size,channel_size,784]
        img = img.view(num_img,  -1)  # 将图片展开为28*28=784
        real_img = img.to(device)
        real_label = label_onehot
        fake_label = torch.zeros((num_img,10)).to(device)
        # print(img.shape)
        # 计算真实图片的损失
        real_out = D(real_img)  # 将真实图片放入判别器中
        d_loss_real = criterion(real_out, real_label)  # 得到真实图片的loss
        # 计算假的图片的损失
        z = torch.randn(num_img, z_dimension+10).to(device)  # 随机生成一些噪声
        fake_img = G(z)  # 随机噪声放入生成网络中,生成一张假的图片
        fake_out = D(fake_img)  # 判别器判断假的图片
        d_loss_fake = criterion(fake_out, fake_label)  # 得到假的图片的loss

        # 损失函数和优化
        d_loss = d_loss_real + d_loss_fake  # 损失包括判真损失和判假损失
        d_optimizer.zero_grad()  # 在反向传播之前,先将梯度归0
        d_loss.backward()  # 将误差反向传播
        d_optimizer.step()  # 更新参数

        # ==================训练生成器============================
        ################################生成网络的训练###############################
        z = torch.randn(num_img, z_dimension).to(device) # 得到随机噪声
        z = torch.cat([z, real_label],1)
        fake_img = G(z)
        output = D(fake_img)
        g_loss = criterion(output, real_label)
        # bp and optimize
        g_optimizer.zero_grad()  # 梯度归0
        g_loss.backward()  # 进行反向传播
        g_optimizer.step()  # .step()一般用在反向传播后面,用于更新生成网络的参数
        # 打印中间的损失
        # try:
        if (i + 1) % 100 == 0:
            print('Epoch[{}/{}],d_loss:{:.6f},g_loss:{:.6f} '.format(
                epoch, num_epoch, d_loss.item(), g_loss.item(),
            ))
        # except BaseException as e:
        #     pass
        if epoch == 0:
            real_images = real_img.cpu().clamp(0,1).view(-1,1,28,28).data
            save_image(real_images, './img_CGAN/real_images.png')
        if i == len(dataloader)-1:
            fake_images = fake_img.cpu().clamp(0,1).view(-1,1,28,28).data
            save_image(fake_images, './img_CGAN/fake_images-{}.png'.format(epoch + 1))
# 保存模型
torch.save(G.state_dict(), './generator_CGAN_z100.pth')
torch.save(D.state_dict(), './discriminator_CGAN_z100.pth')

DCGAN

# coding=utf-8
import torch.autograd
import torch.nn as nn
from torchvision import transforms
from torchvision import datasets
from torchvision.utils import save_image
import os

# 创建文件夹
if not os.path.exists('./img_DCGAN'):
    os.mkdir('./img_DCGAN')
# GPU
device = 'cuda' if torch.cuda.is_available() else 'cpu'
def to_img(x):
    out = 0.5 * (x + 1)
    out = out.clamp(0, 1)  # Clamp函数可以将随机变化的数值限制在一个给定的区间[min, max]内:
    out = out.view(-1, 1, 28, 28)  # view()函数作用是将一个多行的Tensor,拼接成一行
    return out

batch_size = 128
num_epoch = 100
z_dimension = 100

# 图形啊处理过程
img_transform = transforms.Compose([
    transforms.ToTensor(),
    # transforms.Lambda(lambda x: x.repeat(3,1,1)),
    transforms.Normalize(mean=[0.5], std=[0.5])
])

# mnist dataset mnist数据集下载
mnist = datasets.MNIST(
    root='./data/', train=True, transform=img_transform, download=True
)

# data loader 数据载入
dataloader = torch.utils.data.DataLoader(
    dataset=mnist, batch_size=batch_size, shuffle=True
)

# 定义判别器  #####Discriminator######使用多层网络来作为判别器
class discriminator(nn.Module):
    def __init__(self):
        super(discriminator, self).__init__()
        self.dis = nn.Sequential(

            nn.Conv2d(1,32,3,stride=1,padding=1),
            nn.LeakyReLU(0.2,True),
            nn.MaxPool2d((2,2)),

            nn.Conv2d(32,64,3,stride=1,padding=1),
            nn.LeakyReLU(0.2,True),
            nn.MaxPool2d((2,2)),
        )
        self.fc = nn.Sequential(
            nn.Linear(7*7*64,1024),
            nn.LeakyReLU(0.2,True),
            nn.Linear(1024,1),
            nn.Sigmoid()
        )

    def forward(self, x):
        x = self.dis(x)
        x=x.view(x.size(0),-1)
        x=self.fc(x)
        return x

####### 定义生成器 Generator #####
class generator(nn.Module):
    def __init__(self,input_size,num_feature):
        super(generator, self).__init__()
        self.fc=nn.Linear(input_size,num_feature)
        self.br=nn.Sequential(
            nn.BatchNorm2d(1),
            nn.ReLU(True),
        )
        self.gen = nn.Sequential(
            nn.Conv2d(1,64,3,stride=1,padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(True),

            nn.Conv2d(64,32,3,stride=1,padding=1),
            nn.BatchNorm2d(32),
            nn.ReLU(True),

            nn.Conv2d(32,1,3,stride=2,padding=1),
            nn.Tanh(),
        )

    def forward(self, x):
        x = self.fc(x)
        x=x.view(x.shape[0],1,56,56)
        x=self.br(x)
        x=self.gen(x)
        return x

# 创建对象
D = discriminator()
G = generator(100,1*56*56)
D = D.to(device)
G = G.to(device)

#########判别器训练train#####################
# 分为两部分:1、真的图像判别为真;2、假的图像判别为假
# 此过程中,生成器参数不断更新

# 首先需要定义loss的度量方式  (二分类的交叉熵)
# 其次定义 优化函数,优化函数的学习率为0.0003
criterion = nn.BCELoss()  # 是单目标二分类交叉熵函数
d_optimizer = torch.optim.Adam(D.parameters(), lr=0.0003)
g_optimizer = torch.optim.Adam(G.parameters(), lr=0.0003)

###########################进入训练##判别器的判断过程#####################

for epoch in range(num_epoch):  # 进行多个epoch的训练
    for i, (img, _) in enumerate(dataloader):
        for a in range(3):
            num_img = img.size(0)
            # view()函数作用把img变成[batch_size,channel_size,784]
            img = img.view(num_img,  1,28,28)  # 将图片展开为28*28=784
            real_img = img.to(device) # 将tensor变成Variable放入计算图中
            real_label = torch.ones(num_img).to(device)# 定义真实的图片label为1
            fake_label = torch.zeros(num_img).to(device)  # 定义假的图片的label为0
            # print(img.shape)
            # 计算真实图片的损失
            real_out = D(real_img)  # 将真实图片放入判别器中
            d_loss_real = criterion(real_out, real_label)  # 得到真实图片的loss
            real_scores = real_out  # 得到真实图片的判别值,输出的值越接近1越好

            # 计算假的图片的损失
            z = torch.randn(num_img, z_dimension).to(device)  # 随机生成一些噪声
            fake_img = G(z)  # 随机噪声放入生成网络中,生成一张假的图片
            fake_out = D(fake_img)  # 判别器判断假的图片
            d_loss_fake = criterion(fake_out, fake_label)  # 得到假的图片的loss
            fake_scores = fake_out  # 得到假图片的判别值,对于判别器来说,假图片的损失越接近0越好

            # 损失函数和优化
            d_loss = d_loss_real + d_loss_fake  # 损失包括判真损失和判假损失
            d_optimizer.zero_grad()  # 在反向传播之前,先将梯度归0
            d_loss.backward()  # 将误差反向传播
            d_optimizer.step()  # 更新参数

        # ==================训练生成器============================
        ################################生成网络的训练###############################
        # 原理:目的是希望生成的假的图片被判别器判断为真的图片,
        # 在此过程中,将判别器固定,将假的图片传入判别器的结果与真实的label对应,
        # 反向传播更新的参数是生成网络里面的参数,
        # 这样可以通过更新生成网络里面的参数,来训练网络,使得生成的图片让判别器以为是真的
        # 这样就达到了对抗的目的

        # 计算假的图片的损失

        z = torch.randn(num_img, z_dimension).to(device) # 得到随机噪声
        fake_img = G(z)  # 随机噪声输入到生成器中,得到一副假的图片
        output = D(fake_img).squeeze(1)  # 经过判别器得到的结果
        g_loss = criterion(output, real_label)  # 得到的假的图片与真实的图片的label的loss

        # bp and optimize
        g_optimizer.zero_grad()  # 梯度归0
        g_loss.backward()  # 进行反向传播
        g_optimizer.step()  # .step()一般用在反向传播后面,用于更新生成网络的参数

        # 打印中间的损失
        # try:
        if (i + 1) % 100 == 0:
            print('Epoch[{}/{}],d_loss:{:.6f},g_loss:{:.6f} '
                  'D real: {:.6f},D fake: {:.6f}'.format(
                epoch, num_epoch, d_loss.item(), g_loss.item(),
                torch.mean(real_scores).item(), torch.mean(fake_scores).item()  # 打印的是真实图片的损失均值
            ))
        # except BaseException as e:
        #     pass

        if epoch == 0:
            real_images = to_img(real_img.cpu().data)
            save_image(real_images, './img_DCGAN/real_images.png')

        fake_images = to_img(fake_img.cpu().data)
        save_image(fake_images, './img_DCGAN/fake_images-{}.png'.format(epoch + 1))
# 保存模型
torch.save(G.state_dict(), './generator_DCGAN.pth')
torch.save(D.state_dict(), './discriminator_DCGAN.pth')

WGAN

# coding=utf-8
import torch.autograd
import torch.nn as nn
from torchvision import transforms
from torchvision import datasets
from torchvision.utils import save_image
import os

# 创建文件夹
if not os.path.exists('./img_WGAN'):
    os.mkdir('./img_WGAN')
# GPU
device = 'cuda' if torch.cuda.is_available() else 'cpu'

def to_img(x):
    out = 0.5 * (x + 1)
    out = out.clamp(0, 1)  # Clamp函数可以将随机变化的数值限制在一个给定的区间[min, max]内:
    out = out.view(-1, 1, 28, 28)  # view()函数作用是将一个多行的Tensor,拼接成一行
    return out

batch_size = 128
num_epoch = 100
z_dimension = 100
c = 0.005
# 图形啊处理过程
img_transform = transforms.Compose([
    transforms.ToTensor(),
    # transforms.Lambda(lambda x: x.repeat(3,1,1)),
    transforms.Normalize(mean=[0.5], std=[0.5])
])

# mnist dataset mnist数据集下载
mnist = datasets.MNIST(
    root='./data/', train=True, transform=img_transform, download=True
)

# data loader 数据载入
dataloader = torch.utils.data.DataLoader(
    dataset=mnist, batch_size=batch_size, shuffle=True
)

# 定义判别器  #####Discriminator######使用多层网络来作为判别器

# 将图片28x28展开成784,然后通过多层感知器,中间经过斜率设置为0.2的LeakyReLU激活函数,
# 最后接sigmoid激活函数得到一个0到1之间的概率进行二分类。
class discriminator(nn.Module):
    def __init__(self):
        super(discriminator, self).__init__()
        self.dis = nn.Sequential(
            nn.Linear(784, 256),  # 输入特征数为784,输出为256
            nn.LeakyReLU(0.2),  # 进行非线性映射
            nn.Linear(256, 256),  # 进行一个线性映射
            nn.LeakyReLU(0.2),
            nn.Linear(256, 1),
            # nn.Sigmoid()  # 也是一个激活函数,二分类问题中,
            # sigmoid可以班实数映射到【0,1】,作为概率值,
            # 多分类用softmax函数
        )

    def forward(self, x):
        x = self.dis(x)
        return x

####### 定义生成器 Generator #####
# 输入一个100维的0~1之间的高斯分布,然后通过第一层线性变换将其映射到256维,
# 然后通过LeakyReLU激活函数,接着进行一个线性变换,再经过一个LeakyReLU激活函数,
# 然后经过线性变换将其变成784维,最后经过Tanh激活函数是希望生成的假的图片数据分布
# 能够在-1~1之间。
class generator(nn.Module):
    def __init__(self):
        super(generator, self).__init__()
        self.gen = nn.Sequential(
            nn.Linear(100, 256),  # 用线性变换将输入映射到256维
            nn.ReLU(True),  # relu激活
            nn.Linear(256, 256),  # 线性变换
            nn.ReLU(True),  # relu激活
            nn.Linear(256, 784),  # 线性变换
            nn.Tanh()  # Tanh激活使得生成数据分布在【-1,1】之间
        )

    def forward(self, x):
        x = self.gen(x)
        return x

# 创建对象
D = discriminator()
G = generator()
G.load_state_dict(torch.load('./generator_WGAN.pth'))
D.load_state_dict(torch.load('./discriminator_WGAN.pth'))
D = D.to(device)
G = G.to(device)

#########判别器训练train#####################
# 分为两部分:1、真的图像判别为真;2、假的图像判别为假
# 此过程中,生成器参数不断更新

# 首先需要定义loss的度量方式  (二分类的交叉熵)
# 其次定义 优化函数,优化函数的学习率为0.0003
# criterion = nn.BCELoss()  # 是单目标二分类交叉熵函数
d_optimizer = torch.optim.Adam(D.parameters(), lr=0.0001)
g_optimizer = torch.optim.Adam(G.parameters(), lr=0.0001)

###########################进入训练##判别器的判断过程#####################

for epoch in range(num_epoch):  # 进行多个epoch的训练
    for i, (img, _) in enumerate(dataloader):
        for a in range(5):
            num_img = img.size(0)
            # view()函数作用把img变成[batch_size,channel_size,784]
            img = img.view(num_img,  -1)  # 将图片展开为28*28=784
            real_img = img.to(device) # 将tensor变成Variable放入计算图中
            real_label = torch.ones(num_img).to(device)  # 定义真实的图片label为1
            fake_label = torch.zeros(num_img).to(device) # 定义假的图片的label为0
            # print(img.shape)
            # 计算真实图片的损失
            real_out = D(real_img)  # 将真实图片放入判别器中
            real_scores = real_out  # 得到真实图片的判别值,输出的值越接近1越好

            # 计算假的图片的损失
            z = torch.randn(num_img, z_dimension).to(device)  # 随机生成一些噪声
            fake_img = G(z)  # 随机噪声放入生成网络中,生成一张假的图片
            fake_out = D(fake_img)  # 判别器判断假的图片
            d_loss = torch.mean(fake_out)-torch.mean(real_out)  # 得到假的图片的loss
            fake_scores = fake_out  # 得到假图片的判别值,对于判别器来说,假图片的损失越接近0越好

            # 损失函数和优化
            # d_loss = d_loss_real + d_loss_fake  # 损失包括判真损失和判假损失
            d_optimizer.zero_grad()  # 在反向传播之前,先将梯度归0
            d_loss.backward()  # 将误差反向传播
            d_optimizer.step()  # 更新参数
            # weight Clipping WGAN
            for layer in D.dis:
                if (layer.__class__.__name__ == 'Linear'):
                    layer.weight.requires_grad = False
                    layer.weight.clamp_(-c, c)
                    layer.weight.requires_grad = True
        # ==================训练生成器============================
        ################################生成网络的训练###############################
        # 原理:目的是希望生成的假的图片被判别器判断为真的图片,
        # 在此过程中,将判别器固定,将假的图片传入判别器的结果与真实的label对应,
        # 反向传播更新的参数是生成网络里面的参数,
        # 这样可以通过更新生成网络里面的参数,来训练网络,使得生成的图片让判别器以为是真的
        # 这样就达到了对抗的目的

        # 计算假的图片的损失

        z = torch.randn(num_img, z_dimension).to(device) # 得到随机噪声
        fake_img = G(z)  # 随机噪声输入到生成器中,得到一副假的图片
        output = D(fake_img).squeeze(1)  # 经过判别器得到的结果
        g_loss = torch.mean(-output)  # 得到的假的图片与真实的图片的label的loss

        # bp and optimize
        g_optimizer.zero_grad()  # 梯度归0
        g_loss.backward()  # 进行反向传播
        g_optimizer.step()  # .step()一般用在反向传播后面,用于更新生成网络的参数

        # 打印中间的损失
        # try:
        if (i + 1) % 100 == 0:
            print('Epoch[{}/{}],d_loss:{:.6f},g_loss:{:.6f} '
                  'D real: {:.6f},D fake: {:.6f}'.format(
                epoch, num_epoch, d_loss.item(), g_loss.item(),
                torch.mean(real_scores).item(), torch.mean(fake_scores).item()  # 打印的是真实图片的损失均值
            ))
        # except BaseException as e:
        #     pass

        if epoch == 0:
            real_images = to_img(real_img.cpu().data)
            save_image(real_images, './img_WGAN/real_images.png')

        fake_images = to_img(fake_img.cpu().data)
        save_image(fake_images, './img_WGAN/fake_images-{}.png'.format(epoch + 501))
# 保存模型
torch.save(G.state_dict(), './generator_WGAN.pth')
torch.save(D.state_dict(), './discriminator_WGAN.pth')

WGAN-GP

import torch.autograd
import torch.nn as nn
from torchvision import transforms
from torchvision import datasets
from torchvision.utils import save_image
import os

# 创建文件夹
if not os.path.exists('./img_WGANGP'):
    os.mkdir('./img_WGANGP')
# GPU
device = 'cuda' if torch.cuda.is_available() else 'cpu'
def to_img(x):
    out = 0.5 * (x + 1)
    out = out.clamp(0, 1)  # Clamp函数可以将随机变化的数值限制在一个给定的区间[min, max]内:
    out = out.view(-1, 1, 28, 28)  # view()函数作用是将一个多行的Tensor,拼接成一行
    return out

batch_size = 128
num_epoch = 100
z_dimension = 100
lambda_=10
# 图形啊处理过程
img_transform = transforms.Compose([
    transforms.ToTensor(),
    # transforms.Lambda(lambda x: x.repeat(3,1,1)),
    transforms.Normalize(mean=[0.5], std=[0.5])
])

# mnist dataset mnist数据集下载
mnist = datasets.MNIST(
    root='./data/', train=True, transform=img_transform, download=True
)

# data loader 数据载入
dataloader = torch.utils.data.DataLoader(
    dataset=mnist, batch_size=batch_size, shuffle=True
)

# 定义判别器  #####Discriminator######使用多层网络来作为判别器

# 将图片28x28展开成784,然后通过多层感知器,中间经过斜率设置为0.2的LeakyReLU激活函数,
# 最后接sigmoid激活函数得到一个0到1之间的概率进行二分类。
class discriminator(nn.Module):
    def __init__(self):
        super(discriminator, self).__init__()
        self.dis = nn.Sequential(
            nn.Linear(784, 256),  # 输入特征数为784,输出为256
            nn.LeakyReLU(0.2),  # 进行非线性映射
            nn.Linear(256, 256),  # 进行一个线性映射
            nn.LeakyReLU(0.2),
            nn.Linear(256, 1),
            # nn.Sigmoid()  # 也是一个激活函数,二分类问题中,
            # sigmoid可以班实数映射到【0,1】,作为概率值,
            # 多分类用softmax函数
        )

    def forward(self, x):
        x = self.dis(x)
        return x

####### 定义生成器 Generator #####
# 输入一个100维的0~1之间的高斯分布,然后通过第一层线性变换将其映射到256维,
# 然后通过LeakyReLU激活函数,接着进行一个线性变换,再经过一个LeakyReLU激活函数,
# 然后经过线性变换将其变成784维,最后经过Tanh激活函数是希望生成的假的图片数据分布
# 能够在-1~1之间。
class generator(nn.Module):
    def __init__(self):
        super(generator, self).__init__()
        self.gen = nn.Sequential(
            nn.Linear(100, 256),  # 用线性变换将输入映射到256维
            nn.ReLU(True),  # relu激活
            nn.Linear(256, 256),  # 线性变换
            nn.ReLU(True),  # relu激活
            nn.Linear(256, 784),  # 线性变换
            nn.Tanh()  # Tanh激活使得生成数据分布在【-1,1】之间
        )

    def forward(self, x):
        x = self.gen(x)
        return x

# 创建对象
D = discriminator()
G = generator()
D = D.to(device)
G = G.to(device)

#########判别器训练train#####################
# 分为两部分:1、真的图像判别为真;2、假的图像判别为假
# 此过程中,生成器参数不断更新

# 首先需要定义loss的度量方式  (二分类的交叉熵)
# 其次定义 优化函数,优化函数的学习率为0.0003
# criterion = nn.BCELoss()  # 是单目标二分类交叉熵函数
d_optimizer = torch.optim.Adam(D.parameters(), lr=0.0001)
g_optimizer = torch.optim.Adam(G.parameters(), lr=0.0001)

###########################进入训练##判别器的判断过程#####################

for epoch in range(num_epoch):  # 进行多个epoch的训练
    for i, (img, _) in enumerate(dataloader):
        for a in range(5):
            num_img = img.size(0)
            # view()函数作用把img变成[batch_size,channel_size,784]
            img = img.view(num_img,  -1)  # 将图片展开为28*28=784
            real_img = img.to(device)  # 将tensor变成Variable放入计算图中

            real_label = torch.ones(num_img).to(device)  # 定义真实的图片label为1
            fake_label = torch.zeros(num_img).to(device) # 定义假的图片的label为0
            # print(img.shape)
            # 计算真实图片的损失
            real_out = D(real_img)  # 将真实图片放入判别器中
            real_scores = real_out  # 得到真实图片的判别值,输出的值越接近1越好

            # 计算假的图片的损失
            z = torch.randn(num_img, z_dimension).to(device)  # 随机生成一些噪声
            fake_img = G(z)  # 随机噪声放入生成网络中,生成一张假的图片

            alpha = torch.rand((num_img, 1, 1, 1)).to(device)
            x_hat = alpha * real_img + (1 - alpha) * fake_img
            pred_hat = D(x_hat)
            gradients = torch.autograd.grad(outputs=pred_hat, inputs=x_hat, grad_outputs=torch.ones(pred_hat.size()).to(device),
                             create_graph=True, retain_graph=True, only_inputs=True)[0]
            gradient_penalty = lambda_*((gradients.view(gradients.size()[0], -1).norm(2, 1) - 1) ** 2).mean()
            fake_out = D(fake_img)  # 判别器判断假的图片
            d_loss = torch.mean(fake_out)-torch.mean(real_out)+gradient_penalty  # 得到假的图片的loss
            fake_scores = fake_out  # 得到假图片的判别值,对于判别器来说,假图片的损失越接近0越好

            # 损失函数和优化
            # d_loss = d_loss_real + d_loss_fake  # 损失包括判真损失和判假损失
            d_optimizer.zero_grad()  # 在反向传播之前,先将梯度归0
            d_loss.backward()  # 将误差反向传播
            d_optimizer.step()  # 更新参数
        # ==================训练生成器============================
        ################################生成网络的训练###############################
        # 原理:目的是希望生成的假的图片被判别器判断为真的图片,
        # 在此过程中,将判别器固定,将假的图片传入判别器的结果与真实的label对应,
        # 反向传播更新的参数是生成网络里面的参数,
        # 这样可以通过更新生成网络里面的参数,来训练网络,使得生成的图片让判别器以为是真的
        # 这样就达到了对抗的目的

        # 计算假的图片的损失

        z = torch.randn(num_img, z_dimension).to(device) # 得到随机噪声
        fake_img = G(z)  # 随机噪声输入到生成器中,得到一副假的图片
        output = D(fake_img).squeeze(1)  # 经过判别器得到的结果
        g_loss = torch.mean(-output)  # 得到的假的图片与真实的图片的label的loss

        # bp and optimize
        g_optimizer.zero_grad()  # 梯度归0
        g_loss.backward()  # 进行反向传播
        g_optimizer.step()  # .step()一般用在反向传播后面,用于更新生成网络的参数

        # 打印中间的损失
        # try:
        if (i + 1) % 100 == 0:
            print('Epoch[{}/{}],d_loss:{:.6f},g_loss:{:.6f} '
                  'D real: {:.6f},D fake: {:.6f}'.format(
                epoch, num_epoch, d_loss.item(), g_loss.item(),
                torch.mean(real_scores).item(), torch.mean(fake_scores).item()  # 打印的是真实图片的损失均值
            ))
        # except BaseException as e:
        #     pass

        if epoch == 0:
            real_images = to_img(real_img.cpu().data)
            save_image(real_images, './img_WGANGP/real_images.png')

        fake_images = to_img(fake_img.cpu().data)
        save_image(fake_images, './img_WGANGP/fake_images-{}.png'.format(epoch + 1))
# 保存模型
torch.save(G.state_dict(), './generator_WGANGP.pth')
torch.save(D.state_dict(), './discriminator_WGANGP.pth')

WDCGAN

import torch.autograd
import torch.nn as nn
from torchvision import transforms
from torchvision import datasets
from torchvision.utils import save_image
import os

# 创建文件夹
if not os.path.exists('./img_WGANGP'):
    os.mkdir('./img_WGANGP')
# GPU
device = 'cuda' if torch.cuda.is_available() else 'cpu'
def to_img(x):
    out = 0.5 * (x + 1)
    out = out.clamp(0, 1)  # Clamp函数可以将随机变化的数值限制在一个给定的区间[min, max]内:
    out = out.view(-1, 1, 28, 28)  # view()函数作用是将一个多行的Tensor,拼接成一行
    return out

batch_size = 128
num_epoch = 100
z_dimension = 100
lambda_=10
# 图形啊处理过程
img_transform = transforms.Compose([
    transforms.ToTensor(),
    # transforms.Lambda(lambda x: x.repeat(3,1,1)),
    transforms.Normalize(mean=[0.5], std=[0.5])
])

# mnist dataset mnist数据集下载
mnist = datasets.MNIST(
    root='./data/', train=True, transform=img_transform, download=True
)

# data loader 数据载入
dataloader = torch.utils.data.DataLoader(
    dataset=mnist, batch_size=batch_size, shuffle=True
)

# 定义判别器  #####Discriminator######使用多层网络来作为判别器

# 将图片28x28展开成784,然后通过多层感知器,中间经过斜率设置为0.2的LeakyReLU激活函数,
# 最后接sigmoid激活函数得到一个0到1之间的概率进行二分类。
class discriminator(nn.Module):
    def __init__(self):
        super(discriminator, self).__init__()
        self.dis = nn.Sequential(
            nn.Linear(784, 256),  # 输入特征数为784,输出为256
            nn.LeakyReLU(0.2),  # 进行非线性映射
            nn.Linear(256, 256),  # 进行一个线性映射
            nn.LeakyReLU(0.2),
            nn.Linear(256, 1),
            # nn.Sigmoid()  # 也是一个激活函数,二分类问题中,
            # sigmoid可以班实数映射到【0,1】,作为概率值,
            # 多分类用softmax函数
        )

    def forward(self, x):
        x = self.dis(x)
        return x

####### 定义生成器 Generator #####
# 输入一个100维的0~1之间的高斯分布,然后通过第一层线性变换将其映射到256维,
# 然后通过LeakyReLU激活函数,接着进行一个线性变换,再经过一个LeakyReLU激活函数,
# 然后经过线性变换将其变成784维,最后经过Tanh激活函数是希望生成的假的图片数据分布
# 能够在-1~1之间。
class generator(nn.Module):
    def __init__(self):
        super(generator, self).__init__()
        self.gen = nn.Sequential(
            nn.Linear(100, 256),  # 用线性变换将输入映射到256维
            nn.ReLU(True),  # relu激活
            nn.Linear(256, 256),  # 线性变换
            nn.ReLU(True),  # relu激活
            nn.Linear(256, 784),  # 线性变换
            nn.Tanh()  # Tanh激活使得生成数据分布在【-1,1】之间
        )

    def forward(self, x):
        x = self.gen(x)
        return x

# 创建对象
D = discriminator()
G = generator()
D = D.to(device)
G = G.to(device)

#########判别器训练train#####################
# 分为两部分:1、真的图像判别为真;2、假的图像判别为假
# 此过程中,生成器参数不断更新

# 首先需要定义loss的度量方式  (二分类的交叉熵)
# 其次定义 优化函数,优化函数的学习率为0.0003
# criterion = nn.BCELoss()  # 是单目标二分类交叉熵函数
d_optimizer = torch.optim.Adam(D.parameters(), lr=0.0001)
g_optimizer = torch.optim.Adam(G.parameters(), lr=0.0001)

###########################进入训练##判别器的判断过程#####################

for epoch in range(num_epoch):  # 进行多个epoch的训练
    for i, (img, _) in enumerate(dataloader):
        for a in range(5):
            num_img = img.size(0)
            # view()函数作用把img变成[batch_size,channel_size,784]
            img = img.view(num_img,  -1)  # 将图片展开为28*28=784
            real_img = img.to(device)  # 将tensor变成Variable放入计算图中

            real_label = torch.ones(num_img).to(device)  # 定义真实的图片label为1
            fake_label = torch.zeros(num_img).to(device) # 定义假的图片的label为0
            # print(img.shape)
            # 计算真实图片的损失
            real_out = D(real_img)  # 将真实图片放入判别器中
            real_scores = real_out  # 得到真实图片的判别值,输出的值越接近1越好

            # 计算假的图片的损失
            z = torch.randn(num_img, z_dimension).to(device)  # 随机生成一些噪声
            fake_img = G(z)  # 随机噪声放入生成网络中,生成一张假的图片

            alpha = torch.rand((num_img, 1, 1, 1)).to(device)
            x_hat = alpha * real_img + (1 - alpha) * fake_img
            pred_hat = D(x_hat)
            gradients = torch.autograd.grad(outputs=pred_hat, inputs=x_hat, grad_outputs=torch.ones(pred_hat.size()).to(device),
                             create_graph=True, retain_graph=True, only_inputs=True)[0]
            gradient_penalty = lambda_*((gradients.view(gradients.size()[0], -1).norm(2, 1) - 1) ** 2).mean()
            fake_out = D(fake_img)  # 判别器判断假的图片
            d_loss = torch.mean(fake_out)-torch.mean(real_out)+gradient_penalty  # 得到假的图片的loss
            fake_scores = fake_out  # 得到假图片的判别值,对于判别器来说,假图片的损失越接近0越好

            # 损失函数和优化
            # d_loss = d_loss_real + d_loss_fake  # 损失包括判真损失和判假损失
            d_optimizer.zero_grad()  # 在反向传播之前,先将梯度归0
            d_loss.backward()  # 将误差反向传播
            d_optimizer.step()  # 更新参数
        # ==================训练生成器============================
        ################################生成网络的训练###############################
        # 原理:目的是希望生成的假的图片被判别器判断为真的图片,
        # 在此过程中,将判别器固定,将假的图片传入判别器的结果与真实的label对应,
        # 反向传播更新的参数是生成网络里面的参数,
        # 这样可以通过更新生成网络里面的参数,来训练网络,使得生成的图片让判别器以为是真的
        # 这样就达到了对抗的目的

        # 计算假的图片的损失

        z = torch.randn(num_img, z_dimension).to(device) # 得到随机噪声
        fake_img = G(z)  # 随机噪声输入到生成器中,得到一副假的图片
        output = D(fake_img).squeeze(1)  # 经过判别器得到的结果
        g_loss = torch.mean(-output)  # 得到的假的图片与真实的图片的label的loss

        # bp and optimize
        g_optimizer.zero_grad()  # 梯度归0
        g_loss.backward()  # 进行反向传播
        g_optimizer.step()  # .step()一般用在反向传播后面,用于更新生成网络的参数

        # 打印中间的损失
        # try:
        if (i + 1) % 100 == 0:
            print('Epoch[{}/{}],d_loss:{:.6f},g_loss:{:.6f} '
                  'D real: {:.6f},D fake: {:.6f}'.format(
                epoch, num_epoch, d_loss.item(), g_loss.item(),
                torch.mean(real_scores).item(), torch.mean(fake_scores).item()  # 打印的是真实图片的损失均值
            ))
        # except BaseException as e:
        #     pass

        if epoch == 0:
            real_images = to_img(real_img.cpu().data)
            save_image(real_images, './img_WGANGP/real_images.png')

        fake_images = to_img(fake_img.cpu().data)
        save_image(fake_images, './img_WGANGP/fake_images-{}.png'.format(epoch + 1))
# 保存模型
torch.save(G.state_dict(), './generator_WGANGP.pth')
torch.save(D.state_dict(), './discriminator_WGANGP.pth')

WDCGAN-GP

import torch.autograd
import torch.nn as nn
from torchvision import transforms
from torchvision import datasets
from torchvision.utils import save_image
import os

# 创建文件夹
if not os.path.exists('./img_WDCGAN-GP'):
    os.mkdir('./img_WDCGAN-GP')
# GPU
device = 'cuda' if torch.cuda.is_available() else 'cpu'
def to_img(x):
    out = 0.5 * (x + 1)
    out = out.clamp(0, 1)  # Clamp函数可以将随机变化的数值限制在一个给定的区间[min, max]内:
    out = out.view(-1, 1, 28, 28)  # view()函数作用是将一个多行的Tensor,拼接成一行
    return out

batch_size = 128
num_epoch = 25
z_dimension = 100
c=0.005
lambda_=10

# 图形啊处理过程
img_transform = transforms.Compose([
    transforms.ToTensor(),
    # transforms.Lambda(lambda x: x.repeat(3,1,1)),
    transforms.Normalize(mean=[0.5], std=[0.5])
])

# mnist dataset mnist数据集下载
mnist = datasets.MNIST(
    root='./data/', train=True, transform=img_transform, download=True
)

# data loader 数据载入
dataloader = torch.utils.data.DataLoader(
    dataset=mnist, batch_size=batch_size, shuffle=True
)

# 定义判别器  #####Discriminator######使用多层网络来作为判别器
class discriminator(nn.Module):
    def __init__(self):
        super(discriminator, self).__init__()
        self.dis = nn.Sequential(
            nn.Conv2d(1,32,3,stride=1,padding=1),
            nn.LeakyReLU(0.2,True),
            nn.MaxPool2d((2,2)),

            nn.Conv2d(32,64,3,stride=1,padding=1),
            nn.LeakyReLU(0.2,True),
            nn.MaxPool2d((2,2)),
        )
        self.fc = nn.Sequential(
            nn.Linear(7*7*64,1024),
            nn.LeakyReLU(0.2,True),
            nn.Linear(1024,1),
            # nn.Sigmoid()
        )

    def forward(self, x):
        x = self.dis(x)
        x=x.view(x.size(0),-1)
        x=self.fc(x)
        return x

####### 定义生成器 Generator #####
class generator(nn.Module):
    def __init__(self,input_size,num_feature):
        super(generator, self).__init__()
        self.fc=nn.Linear(input_size,num_feature)
        self.br=nn.Sequential(
            nn.BatchNorm2d(1),
            nn.ReLU(True),
        )
        self.gen = nn.Sequential(
            nn.Conv2d(1,64,3,stride=1,padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(True),

            nn.Conv2d(64,32,3,stride=1,padding=1),
            nn.BatchNorm2d(32),
            nn.ReLU(True),

            nn.Conv2d(32,1,3,stride=2,padding=1),
            nn.Tanh(),
        )

    def forward(self, x):
        x = self.fc(x)
        x=x.view(x.shape[0],1,56,56)
        x=self.br(x)
        x=self.gen(x)
        return x

# 创建对象
D = discriminator()
G = generator(100,1*56*56)
D = D.to(device)
G = G.to(device)

#########判别器训练train#####################
# 分为两部分:1、真的图像判别为真;2、假的图像判别为假
# 此过程中,生成器参数不断更新

# 首先需要定义loss的度量方式  (二分类的交叉熵)
# 其次定义 优化函数,优化函数的学习率为0.0003
criterion = nn.BCELoss()  # 是单目标二分类交叉熵函数
d_optimizer = torch.optim.Adam(D.parameters(), lr=0.0003)
g_optimizer = torch.optim.Adam(G.parameters(), lr=0.0003)

###########################进入训练##判别器的判断过程#####################

for epoch in range(num_epoch):  # 进行多个epoch的训练
    for i, (img, _) in enumerate(dataloader):
        for a in range(3):
            num_img = img.size(0)
            # view()函数作用把img变成[batch_size,channel_size,784]
            img = img.view(num_img,  1,28,28)  # 将图片展开为28*28=784
            real_img = img.to(device) # 将tensor变成Variable放入计算图中
            real_label = torch.ones(num_img).to(device)# 定义真实的图片label为1
            fake_label = torch.zeros(num_img).to(device)  # 定义假的图片的label为0
            # print(img.shape)
            # 计算真实图片的损失
            real_out = D(real_img)  # 将真实图片放入判别器中
            # d_loss_real = criterion(real_out, real_label)  # 得到真实图片的loss
            real_scores = real_out  # 得到真实图片的判别值,输出的值越接近1越好

            # 计算假的图片的损失
            z = torch.randn(num_img, z_dimension).to(device)  # 随机生成一些噪声
            fake_img = G(z)  # 随机噪声放入生成网络中,生成一张假的图片
            fake_out = D(fake_img)  # 判别器判断假的图片
            # d_loss_fake = criterion(fake_out, fake_label)  # 得到假的图片的loss
            fake_scores = fake_out  # 得到假图片的判别值,对于判别器来说,假图片的损失越接近0越好

            # 损失函数和优化
            alpha = torch.rand((num_img, 1, 1, 1)).to(device)
            x_hat = alpha * real_img + (1 - alpha) * fake_img
            pred_hat = D(x_hat)
            gradients = \
            torch.autograd.grad(outputs=pred_hat, inputs=x_hat, grad_outputs=torch.ones(pred_hat.size()).to(device),
                                create_graph=True, retain_graph=True, only_inputs=True)[0]
            gradient_penalty = lambda_ * ((gradients.view(gradients.size()[0], -1).norm(2, 1) - 1) ** 2).mean()
            d_loss = torch.mean(fake_out)-torch.mean(real_out)+gradient_penalty # 损失包括判真损失和判假损失
            d_optimizer.zero_grad()  # 在反向传播之前,先将梯度归0
            d_loss.backward()  # 将误差反向传播
            d_optimizer.step()  # 更新参数
        # ==================训练生成器============================
        ################################生成网络的训练###############################
        # 原理:目的是希望生成的假的图片被判别器判断为真的图片,
        # 在此过程中,将判别器固定,将假的图片传入判别器的结果与真实的label对应,
        # 反向传播更新的参数是生成网络里面的参数,
        # 这样可以通过更新生成网络里面的参数,来训练网络,使得生成的图片让判别器以为是真的
        # 这样就达到了对抗的目的

        # 计算假的图片的损失

        z = torch.randn(num_img, z_dimension).to(device) # 得到随机噪声
        fake_img = G(z)  # 随机噪声输入到生成器中,得到一副假的图片
        output = D(fake_img).squeeze(1)  # 经过判别器得到的结果
        g_loss = torch.mean(-output)   # 得到的假的图片与真实的图片的label的loss

        # bp and optimize
        g_optimizer.zero_grad()  # 梯度归0
        g_loss.backward()  # 进行反向传播
        g_optimizer.step()  # .step()一般用在反向传播后面,用于更新生成网络的参数

        # 打印中间的损失
        # try:
        if (i + 1) % 100 == 0:
            print('Epoch[{}/{}],d_loss:{:.6f},g_loss:{:.6f} '
                  'D real: {:.6f},D fake: {:.6f}'.format(
                epoch, num_epoch, d_loss.item(), g_loss.item(),
                torch.mean(real_scores).item(), torch.mean(fake_scores).item()  # 打印的是真实图片的损失均值
            ))
        # except BaseException as e:
        #     pass

        if epoch == 0:
            real_images = to_img(real_img.cpu().data)
            save_image(real_images, './img_WDCGAN-GP/real_images.png')

        fake_images = to_img(fake_img.cpu().data)
        save_image(fake_images, './img_WDCGAN-GP/fake_images-{}.png'.format(epoch + 1))
# 保存模型
torch.save(G.state_dict(), './generator_WDCGAN-GP.pth')
torch.save(D.state_dict(), './discriminator_WDCGAN-GP.pth')

VAE-GAN

from __future__ import print_function
import argparse
import os
import random
import math
import torch
import torch.nn as nn
import torch.nn.parallel
import torch.backends.cudnn as cudnn
import torch.optim as optim
import torch.utils.data
import torchvision.datasets as dset
import torchvision.transforms as transforms
from torchvision.utils import save_image
from torchvision.utils import make_grid
import torch.nn.functional as F
import os
# custom weights initialization called on netG and netD
# def weights_init(m):
#     classname = m.__class__.__name__
#     if classname.find('Conv') != -1:
#         m.weight.data.normal_(0.0, 0.02)
#     elif classname.find('BatchNorm') != -1:
#         m.weight.data.normal_(1.0, 0.02)
#         m.bias.data.fill_(0)
#     print(classname)

class VAE(nn.Module):
    def __init__(self):
        super(VAE, self).__init__()
        # 定义编码器
        self.encoder = nn.Sequential(
            nn.Conv2d(1,16,kernel_size=3,stride=2,padding=1),
            nn.BatchNorm2d(16),
            nn.LeakyReLU(0.2,inplace=True),
            nn.Conv2d(16,32,kernel_size=3,stride=2,padding=1),
            nn.BatchNorm2d(32),
            nn.LeakyReLU(0.2,inplace=True),
            nn.Conv2d(32,32,kernel_size=3,stride=1,padding=1),
            nn.BatchNorm2d(32),
            nn.LeakyReLU(0.2,inplace=True),
        )
        self.encoder_fc1=nn.Linear(32*7*7,nz)
        self.encoder_fc2=nn.Linear(32*7*7,nz)
        self.Sigmoid = nn.Sigmoid()
        self.decoder_fc = nn.Linear(nz,32 * 7 * 7)
        self.decoder = nn.Sequential(
            nn.ConvTranspose2d(32, 16, 4, 2, 1),
            nn.ReLU(inplace=True),
            nn.ConvTranspose2d(16, 1, 4, 2, 1),
            nn.Sigmoid(),
        )

    def noise_reparameterize(self,mean,logvar):
        eps = torch.randn(mean.shape).to(device)
        z = mean + eps * torch.exp(logvar)
        return z

    def forward(self, x):
        out1,out2 = self.encoder(x),self.encoder(x)
        mean = self.encoder_fc1(out1.view(out1.shape[0],-1))
        logstd = self.encoder_fc2(out2.view(out2.shape[0],-1))
        z = self.noise_reparameterize(mean,logstd)
        out3 = self.decoder_fc(z)
        out3 = out3.view(out3.shape[0],32,7,7)
        out3 = self.decoder(out3)
        return out3,mean,logstd

class Discriminator(nn.Module):
    def __init__(self):
        super(Discriminator, self).__init__()
        self.dis = nn.Sequential(
            nn.Conv2d(1, 32, 3, stride=1, padding=1),
            nn.LeakyReLU(0.2, True),
            nn.MaxPool2d((2, 2)),

            nn.Conv2d(32, 64, 3, stride=1, padding=1),
            nn.LeakyReLU(0.2, True),
            nn.MaxPool2d((2, 2)),
        )
        self.fc = nn.Sequential(
            nn.Linear(7 * 7 * 64, 1024),
            nn.LeakyReLU(0.2, True),
            nn.Linear(1024, 1),
            nn.Sigmoid()
        )

    def forward(self, input):
        x = self.dis(input)
        x = x.view(x.size(0), -1)
        x = self.fc(x)
        return x.squeeze(1)

def loss_function(recon_x,x,mean,logstd):
    # BCE = F.binary_cross_entropy(recon_x,x,reduction='sum')
    MSE = MSECriterion(recon_x,x)
    # 因为var是标准差的自然对数,先求自然对数然后平方转换成方差
    var = torch.pow(torch.exp(logstd),2)
    KLD = -0.5 * torch.sum(1+torch.log(var)-torch.pow(mean,2)-var)
    return MSE+KLD

if __name__ == '__main__':
    dataset = 'cifar10'
    dataset = 'mnist'
    batchSize = 128
    imageSize = 28
    nz=100
    nepoch=20
    if not os.path.exists('./img_VAE-GAN'):
        os.mkdir('./img_VAE-GAN')
    print("Random Seed: 88")
    random.seed(88)
    torch.manual_seed(88)
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    # 可以优化运行效率
    cudnn.benchmark = True
    if dataset == 'cifar10':
        dataset = dset.CIFAR10(root='./data', download=True,
                               transform=transforms.Compose([
                                   transforms.Resize(64),
                                   transforms.ToTensor(),
                                   transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
                               ])
                               )
        n_channel = 3
    elif dataset=='mnist':
        dataset = dset.MNIST(root='./data',
                             train=True,
                             transform=transforms.Compose([transforms.ToTensor()]),
                             download=True
                             )
        n_channel = 1
    dataloader = torch.utils.data.DataLoader(dataset,
                                             batch_size=batchSize,
                                             shuffle=True)

    print("=====> 构建VAE")
    vae = VAE().to(device)
    # vae.load_state_dict(torch.load('./VAE-GAN-VAE.pth'))
    print("=====> 构建D")
    D = Discriminator().to(device)
    # D.load_state_dict(torch.load('./VAE-GAN-Discriminator.pth'))
    criterion = nn.BCELoss().to(device)
    MSECriterion = nn.MSELoss().to(device)

    print("=====> Setup optimizer")
    optimizerD = optim.Adam(D.parameters(), lr=0.0001)
    optimizerVAE = optim.Adam(vae.parameters(), lr=0.0001)

    gen_win = None
    rec_win = None
    print("=====> Begin training")
    for epoch in range(nepoch):
        for i, (data,label) in enumerate(dataloader, 0):
            ###################################################################
            # (1) Update D network: maximize log(D(x)) + log(1 - D(G(z)))
            ###################################################################
            # train with real
            D.zero_grad()
            data = data.to(device)
            label = label.to(device)
            batch_size = data.shape[0]
            output = D(data)
            # print("output size of netDs: ", output.size())
            real_label = torch.ones(batch_size).to(device)  # 定义真实的图片label为1
            fake_label = torch.zeros(batch_size).to(device)  # 定义假的图片的label为0
            errD_real = criterion(output, real_label)
            errD_real.backward()
            real_data_score = output.mean().item()
            # train with fake, taking the noise vector z as the input of D network
            # 随机产生一个潜在变量,然后通过decoder 产生生成图片
            z = torch.randn(batch_size, nz).to(device)
            # 通过vae的decoder把潜在变量z变成虚假图片
            fake_data = vae.decoder_fc(z).view(z.shape[0], 32, 7, 7)
            fake_data = vae.decoder(fake_data)
            output = D(fake_data)
            errD_fake = criterion(output, fake_label)
            errD_fake.backward()
            # fake_data_score用来输出查看的,是虚假照片的评分,0最假,1为真
            fake_data_score = output.data.mean()
            errD = errD_real + errD_fake
            optimizerD.step()
            ###################################################
            # (2) Update G network which is the decoder of VAE
            ###################################################
            recon_data,mean,logstd = vae(data)
            vae.zero_grad()
            vae_loss = loss_function(recon_data,data,mean,logstd)
            vae_loss.backward(retain_graph=True)
            optimizerVAE.step()
            ###############################################
            # (3) Update G network: maximize log(D(G(z)))
            ###############################################
            vae.zero_grad()
            real_label = torch.ones(batch_size).to(device)  # 定义真实的图片label为1
            output = D(recon_data)
            errVAE = criterion(output, real_label)
            errVAE.backward()
            D_G_z2 = output.mean().item()
            optimizerVAE.step()
            if i%100==0:
                print('[%d/%d][%d/%d] real_score: %.4f fake_score: %.4f '
                      'Loss_D: %.4f '
                      % (epoch, nepoch, i, len(dataloader),
                         real_data_score,
                         fake_data_score,
                         errD.item()))
            if epoch==0:
                real_images = make_grid(data.cpu(), nrow=8, normalize=True).detach()
                save_image(real_images, './img_VAE-GAN/real_images.png')
            sample = torch.randn(80, nz).to(device)
            output = vae.decoder_fc(sample)
            output = vae.decoder(output.view(output.shape[0], 32, 7, 7))
            fake_images = make_grid(output.cpu(), nrow=8, normalize=True).detach()
            save_image(fake_images, './img_VAE-GAN/fake_images-{}.png'.format(epoch + 1))
torch.save(vae.state_dict(), './VAE-GAN-VAE_epoch5.pth')
torch.save(D.state_dict(),'./VAE-GAN-Discriminator_epoch5.pth')


CVAE-GAN

import random
import torch
import torch.nn as nn
import torch.nn.parallel
import torch.backends.cudnn as cudnn
import torch.optim as optim
import torch.utils.data
import torchvision.datasets as dset
import torchvision.transforms as transforms
from torchvision.utils import save_image
from torchvision.utils import make_grid
import os
class VAE(nn.Module):
    def __init__(self):
        super(VAE, self).__init__()
        # 定义编码器
        self.encoder_conv = nn.Sequential(
            nn.Conv2d(1,16,kernel_size=3,stride=2,padding=1),
            nn.BatchNorm2d(16),
            nn.LeakyReLU(0.2,inplace=True),
            nn.Conv2d(16,32,kernel_size=3,stride=2,padding=1),
            nn.BatchNorm2d(32),
            nn.LeakyReLU(0.2,inplace=True),
            nn.Conv2d(32,32,kernel_size=3,stride=1,padding=1),
            nn.BatchNorm2d(32),
            nn.LeakyReLU(0.2,inplace=True),
        )
        self.encoder_fc1=nn.Linear(32*7*7,nz)
        self.encoder_fc2=nn.Linear(32*7*7,nz)
        self.Sigmoid = nn.Sigmoid()
        self.decoder_fc = nn.Linear(nz+10,32 * 7 * 7)
        self.decoder_deconv = nn.Sequential(
            nn.ConvTranspose2d(32, 16, 4, 2, 1),
            nn.ReLU(inplace=True),
            nn.ConvTranspose2d(16, 1, 4, 2, 1),
            nn.Sigmoid(),
        )

    def noise_reparameterize(self,mean,logvar):
        eps = torch.randn(mean.shape).to(device)
        z = mean + eps * torch.exp(logvar)
        return z

    def forward(self, x):
        z = self.encoder(x)
        output = self.decoder(z)
        return output
    def encoder(self,x):
        out1, out2 = self.encoder_conv(x), self.encoder_conv(x)
        mean = self.encoder_fc1(out1.view(out1.shape[0], -1))
        logstd = self.encoder_fc2(out2.view(out2.shape[0], -1))
        z = self.noise_reparameterize(mean, logstd)
        return z,mean,logstd
    def decoder(self,z):
        out3 = self.decoder_fc(z)
        out3 = out3.view(out3.shape[0], 32, 7, 7)
        out3 = self.decoder_deconv(out3)
        return out3
class Discriminator(nn.Module):
    def __init__(self,outputn=1):
        super(Discriminator, self).__init__()
        self.dis = nn.Sequential(
            nn.Conv2d(1, 32, 3, stride=1, padding=1),
            nn.LeakyReLU(0.2, True),
            nn.MaxPool2d((2, 2)),

            nn.Conv2d(32, 64, 3, stride=1, padding=1),
            nn.LeakyReLU(0.2, True),
            nn.MaxPool2d((2, 2)),
        )
        self.fc = nn.Sequential(
            nn.Linear(7 * 7 * 64, 1024),
            nn.LeakyReLU(0.2, True),
            nn.Linear(1024, outputn),
            nn.Sigmoid()
        )

    def forward(self, input):
        x = self.dis(input)
        x = x.view(x.size(0), -1)
        x = self.fc(x)
        return x.squeeze(1)

def loss_function(recon_x,x,mean,logstd):
    # BCE = F.binary_cross_entropy(recon_x,x,reduction='sum')
    MSE = MSECriterion(recon_x,x)
    # 因为var是标准差的自然对数,先求自然对数然后平方转换成方差
    var = torch.pow(torch.exp(logstd),2)
    KLD = -0.5 * torch.sum(1+torch.log(var)-torch.pow(mean,2)-var)
    return MSE+KLD

if __name__ == '__main__':
    dataset = 'cifar10'
    dataset = 'mnist'
    batchSize = 128
    imageSize = 28
    nz=100
    nepoch=20
    if not os.path.exists('./img_CVAE-GAN'):
        os.mkdir('./img_CVAE-GAN')
    print("Random Seed: 88")
    random.seed(88)
    torch.manual_seed(88)
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    # 可以优化运行效率
    cudnn.benchmark = True
    dataset = dset.MNIST(root='./data',
                         train=True,
                         transform=transforms.Compose([transforms.ToTensor()]),
                         download=True
                         )
    n_channel = 1
    dataloader = torch.utils.data.DataLoader(dataset,
                                             batch_size=batchSize,
                                             shuffle=True)

    print("=====> 构建VAE")
    vae = VAE().to(device)
    vae.load_state_dict(torch.load('./CVAE-GAN-VAE.pth'))
    print("=====> 构建D")
    D = Discriminator(1).to(device)
    D.load_state_dict(torch.load('./CVAE-GAN-Discriminator.pth'))
    print("=====> 构建C")
    C = Discriminator(10).to(device)
    C.load_state_dict(torch.load('./CVAE-GAN-Classifier.pth'))
    criterion = nn.BCELoss().to(device)
    MSECriterion = nn.MSELoss().to(device)

    print("=====> Setup optimizer")
    optimizerD = optim.Adam(D.parameters(), lr=0.0001)
    optimizerC = optim.Adam(C.parameters(), lr=0.0001)
    optimizerVAE = optim.Adam(vae.parameters(), lr=0.0001)

    for epoch in range(nepoch):
        for i, (data,label) in enumerate(dataloader, 0):
            # 先处理一下数据
            data = data.to(device)
            label_onehot = torch.zeros((data.shape[0], 10)).to(device)
            label_onehot[torch.arange(data.shape[0]), label] = 1
            batch_size = data.shape[0]
            # 先训练C
            output = C(data)
            real_label = label_onehot.to(device)  # 定义真实的图片label为1
            errC = criterion(output, real_label)
            C.zero_grad()
            errC.backward()
            optimizerC.step()
            # 再训练D
            output = D(data)
            real_label = torch.ones(batch_size).to(device)   # 定义真实的图片label为1
            fake_label = torch.zeros(batch_size).to(device)  # 定义假的图片的label为0
            errD_real = criterion(output, real_label)

            z = torch.randn(batch_size, nz + 10).to(device)
            fake_data = vae.decoder(z)
            output = D(fake_data)
            errD_fake = criterion(output, fake_label)

            errD = errD_real+errD_fake
            D.zero_grad()
            errD.backward()
            optimizerD.step()
            # 更新VAE(G)1
            z,mean,logstd = vae.encoder(data)
            z = torch.cat([z,label_onehot],1)
            recon_data = vae.decoder(z)
            vae_loss1 = loss_function(recon_data,data,mean,logstd)
            # 更新VAE(G)2
            output = D(recon_data)
            real_label = torch.ones(batch_size).to(device)
            vae_loss2 = criterion(output,real_label)
            # 更新VAE(G)3
            output = C(recon_data)
            real_label = label_onehot
            vae_loss3 = criterion(output, real_label)

            vae.zero_grad()
            vae_loss = vae_loss1+vae_loss2+vae_loss3
            vae_loss.backward()
            optimizerVAE.step()
            if i%100==0:
                print('[%d/%d][%d/%d] Loss_D: %.4f Loss_C: %.4f Loss_G: %.4f'
                      % (epoch, nepoch, i, len(dataloader),
                         errD.item(),errC.item(),vae_loss.item()))
            if epoch==0:
                real_images = make_grid(data.cpu(), nrow=8, normalize=True).detach()
                save_image(real_images, './img_CVAE-GAN/real_images.png')
            if i == len(dataloader)-1:
                sample = torch.randn(data.shape[0], nz).to(device)
                print(label)
                sample = torch.cat([sample,real_label],1)
                output = vae.decoder(sample)
                fake_images = make_grid(output.cpu(), nrow=8, normalize=True).detach()
                save_image(fake_images, './img_CVAE-GAN/fake_images-{}.png'.format(epoch + 26))
torch.save(vae.state_dict(), './CVAE-GAN-VAE.pth')
torch.save(D.state_dict(),'./CVAE-GAN-Discriminator.pth')
torch.save(C.state_dict(),'./CVAE-GAN-Classifier.pth')

VAE-WGANGP

from __future__ import print_function
import argparse
import os
import random
import math
import torch
import torch.nn as nn
import torch.nn.parallel
import torch.backends.cudnn as cudnn
import torch.optim as optim
import torch.utils.data
import torchvision.datasets as dset
import torchvision.transforms as transforms
from torchvision.utils import save_image
from torchvision.utils import make_grid
import torch.nn.functional as F
import os
class VAE(nn.Module):
    def __init__(self):
        super(VAE, self).__init__()
        # 定义编码器
        self.encoder = nn.Sequential(
            nn.Conv2d(1,16,kernel_size=3,stride=2,padding=1),
            nn.BatchNorm2d(16),
            nn.LeakyReLU(0.2,inplace=True),
            nn.Conv2d(16,32,kernel_size=3,stride=2,padding=1),
            nn.BatchNorm2d(32),
            nn.LeakyReLU(0.2,inplace=True),
            nn.Conv2d(32,32,kernel_size=3,stride=1,padding=1),
            nn.BatchNorm2d(32),
            nn.LeakyReLU(0.2,inplace=True),
        )
        self.encoder_fc1=nn.Linear(32*7*7,nz)
        self.encoder_fc2=nn.Linear(32*7*7,nz)
        self.Sigmoid = nn.Sigmoid()
        self.decoder_fc = nn.Linear(nz,32 * 7 * 7)
        self.decoder = nn.Sequential(
            nn.ConvTranspose2d(32, 16, 4, 2, 1),
            nn.ReLU(inplace=True),
            nn.ConvTranspose2d(16, 1, 4, 2, 1),
            nn.Sigmoid(),
        )

    def noise_reparameterize(self,mean,logvar):
        eps = torch.randn(mean.shape).to(device)
        z = mean + eps * torch.exp(logvar)
        return z

    def forward(self, x):
        out1,out2 = self.encoder(x),self.encoder(x)
        mean = self.encoder_fc1(out1.view(out1.shape[0],-1))
        logstd = self.encoder_fc2(out2.view(out2.shape[0],-1))
        z = self.noise_reparameterize(mean,logstd)
        out3 = self.decoder_fc(z)
        out3 = out3.view(out3.shape[0],32,7,7)
        out3 = self.decoder(out3)
        return out3,mean,logstd

class Discriminator(nn.Module):
    def __init__(self):
        super(Discriminator, self).__init__()
        self.dis = nn.Sequential(
            nn.Conv2d(1, 32, 3, stride=1, padding=1),
            nn.LeakyReLU(0.2, True),
            nn.MaxPool2d((2, 2)),

            nn.Conv2d(32, 64, 3, stride=1, padding=1),
            nn.LeakyReLU(0.2, True),
            nn.MaxPool2d((2, 2)),
        )
        self.fc = nn.Sequential(
            nn.Linear(7 * 7 * 64, 1024),
            nn.LeakyReLU(0.2, True),
            nn.Linear(1024, 1),
            # nn.Sigmoid()
        )

    def forward(self, input):
        x = self.dis(input)
        x = x.view(x.size(0), -1)
        x = self.fc(x)
        return x.squeeze(1)

def loss_function(recon_x,x,mean,logstd):
    # BCE = F.binary_cross_entropy(recon_x,x,reduction='sum')
    MSE = MSECriterion(recon_x,x)
    # 因为var是标准差的自然对数,先求自然对数然后平方转换成方差
    var = torch.pow(torch.exp(logstd),2)
    KLD = -0.5 * torch.sum(1+torch.log(var)-torch.pow(mean,2)-var)
    return MSE+KLD

if __name__ == '__main__':
    dataset = 'cifar10'
    dataset = 'mnist'
    batchSize = 128
    imageSize = 28
    nz=100
    nepoch=25
    lambda_ = 10
    if not os.path.exists('./img_VAE-WGANGP'):
        os.mkdir('./img_VAE-WGANGP')
    print("Random Seed: 88")
    random.seed(88)
    torch.manual_seed(88)
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    # 可以优化运行效率
    cudnn.benchmark = True
    if dataset == 'cifar10':
        dataset = dset.CIFAR10(root='./data', download=True,
                               transform=transforms.Compose([
                                   transforms.Resize(64),
                                   transforms.ToTensor(),
                                   transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
                               ])
                               )
        n_channel = 3
    elif dataset=='mnist':
        dataset = dset.MNIST(root='./data',
                             train=True,
                             transform=transforms.Compose([transforms.ToTensor()]),
                             download=True
                             )
        n_channel = 1
    dataloader = torch.utils.data.DataLoader(dataset,
                                             batch_size=batchSize,
                                             shuffle=True)

    print("=====> 构建VAE")
    vae = VAE().to(device)
    # vae.load_state_dict(torch.load('./VAE-WGANGP-VAE_v2.pth'))
    print("=====> 构建D")
    D = Discriminator().to(device)
    # D.load_state_dict(torch.load('./VAE-GAN-Discriminator.pth'))
    criterion = nn.BCELoss().to(device)
    MSECriterion = nn.MSELoss().to(device)

    print("=====> Setup optimizer")
    optimizerD = optim.Adam(D.parameters(), lr=0.0001)
    optimizerVAE = optim.Adam(vae.parameters(), lr=0.0001)

    gen_win = None
    rec_win = None
    print("=====> Begin training")
    for epoch in range(nepoch):
        if(epoch%5==0):
            lambda_ -= 1
        for i, (data,label) in enumerate(dataloader, 0):
            for n in range(1):
                ###################################################################
                # (1) Update D network: maximize log(D(x)) + log(1 - D(G(z)))
                ###################################################################
                # train with real
                D.zero_grad()
                data = data.to(device)
                label = label.to(device)
                batch_size = data.shape[0]
                real_out = D(data)
                # print("output size of netDs: ", output.size())
                real_label = torch.ones(batch_size).to(device)  # 定义真实的图片label为1
                fake_label = torch.zeros(batch_size).to(device)  # 定义假的图片的label为0
                real_data_score = real_out.mean().item()
                # train with fake, taking the noise vector z as the input of D network
                # 随机产生一个潜在变量,然后通过decoder 产生生成图片
                z = torch.randn(batch_size, nz).to(device)
                # 通过vae的decoder把潜在变量z变成虚假图片
                fake_data = vae.decoder_fc(z).view(z.shape[0], 32, 7, 7)
                fake_data = vae.decoder(fake_data)
                fake_out = D(fake_data)
                # fake_data_score用来输出查看的,是虚假照片的评分,0最假,1为真
                fake_data_score = fake_out.mean().item()

                alpha = torch.rand((batch_size, 1, 1, 1)).to(device)
                x_hat = alpha * data + (1 - alpha) * fake_data
                pred_hat = D(x_hat)
                gradients = \
                    torch.autograd.grad(outputs=pred_hat, inputs=x_hat, grad_outputs=torch.ones(pred_hat.size()).to(device),
                                        create_graph=True, retain_graph=True, only_inputs=True)[0]
                gradient_penalty = lambda_ * ((gradients.view(gradients.size()[0], -1).norm(2, 1) - 1) ** 2).mean()
                d_loss = torch.mean(fake_out) - torch.mean(real_out) + gradient_penalty
                d_loss.backward()
                optimizerD.step()
            ###################################################
            # (2) Update G network which is the decoder of VAE
            ###################################################
            recon_data,mean,logstd = vae(data)
            vae.zero_grad()
            vae_loss = loss_function(recon_data,data,mean,logstd)
            vae_loss.backward(retain_graph=True)
            optimizerVAE.step()
            ###############################################
            # (3) Update G network: maximize log(D(G(z)))
            ###############################################
            vae.zero_grad()
            real_label = torch.ones(batch_size).to(device)  # 定义真实的图片label为1
            output = D(recon_data)
            errVAE = torch.mean(-output)
            errVAE.backward()
            D_G_z2 = output.mean().item()
            optimizerVAE.step()
            if i%100==0:
                print('[%d/%d][%d/%d] real_score: %.4f fake_score: %.4f '
                      % (epoch, nepoch, i, len(dataloader),
                         real_data_score,
                         fake_data_score,
                         ))
            if epoch==0:
                real_images = make_grid(data.cpu(), nrow=8, normalize=True).detach()
                save_image(real_images, './img_VAE-WGANGP/real_images.png')
            sample = torch.randn(80, nz).to(device)
            output = vae.decoder_fc(sample)
            output = vae.decoder(output.view(output.shape[0], 32, 7, 7))
            fake_images = make_grid(output.cpu(), nrow=8, normalize=True).detach()
            save_image(fake_images, './img_VAE-WGANGP/fake_images-{}.png'.format(epoch + 201))
    torch.save(vae.state_dict(), './VAE-WGANGP-VAE.pth')
    torch.save(D.state_dict(),'./VAE-WGANGP-Discriminator.pth')

  • 0
    点赞
  • 0
    收藏
    觉得还不错? 一键收藏
  • 0
    评论
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值