pytorch学习之路——11

Generative Adversarial Network

GAN的思想是，您有两个网络，一个是生成器G，另一个是鉴别器D，彼此竞争。 生成器使“伪”数据传递到鉴别器。 鉴别器还可以看到真实的训练数据，并预测接收到的数据是真实的还是伪造的。

 

鉴别器

鉴别器网络将是一个非常典型的线性分类器。 为了使该网络成为通用函数逼近器，我们至少需要一个隐藏层，并且这些隐藏层应具有一个关键属性：所有隐藏层都将向其输出应用Leaky ReLu激活函数。

我们还将采用在输出上使用数值上更稳定的损失函数的方法。 回想一下，我们希望判别器输出值0-1，指示图像是真实的还是伪造的。

 

生成器

生成器网络将与鉴别器网络几乎完全相同，除了我们将tanh激活函数应用于输出层。

由于tanh函数值范围在[-1，1]，所以训练鉴别器时，我们还必须将实际输入图像缩放为像素值介于-1和1之间。

 

训练

1、鉴别器训练

（1）计算真实训练图像上的鉴别器损失
（2）生成假图片
（3）计算伪造的图像上的鉴别器损失
（4）加上真假损失
（5）执行反向传播+优化步骤以更新鉴别器的权重

2、生成器训练

（1）生成假图片
（2）使用翻转的标签计算假图片上的鉴别器损失！
（3）执行反向传播+优化步骤以更新发生器的权重

 

import torch
from torch import nn,optim
import torch.nn.functional as F
from torchvision import datasets,transforms
from torch.utils.data.dataloader import DataLoader

import numpy as np

def get_dataLoader(batch_size):
    transform = transforms.Compose([transforms.ToTensor()])
    trainSet = datasets.MNIST('MNIST_data/',train=True,transform=transform)
    testSet = datasets.MNIST('MNIST_data/',train=False,transform=transform)
    trainLoader = DataLoader(trainSet,batch_size=batch_size)
    testLoader = DataLoader(testSet,batch_size=batch_size)

    return trainLoader,testLoader


class Discriminator(nn.Module):
    def __init__(self,fea_size,hid_size,out_size):
        super().__init__()
        self.in_size = fea_size
        self.fc1 = nn.Linear(fea_size,4*hid_size)
        self.fc2 = nn.Linear(4*hid_size,2*hid_size)
        self.fc3 = nn.Linear(2*hid_size,hid_size)
        self.fc4 = nn.Linear(hid_size,out_size)
        self.drop = nn.Dropout(p=0.3)

    def forward(self,x):
        x = x.view(-1,self.in_size)
        x = F.leaky_relu(self.fc1(x),0.2)
        x = self.drop(x)
        x = F.leaky_relu(self.fc2(x),0.2)
        x = self.drop(x)
        x = F.leaky_relu(self.fc3(x),0.2)
        x = self.drop(x)
        x = self.fc4(x)

        return x


class Generator(nn.Module):
    def __init__(self,lat_size,hid_size,fea_size):
        super().__init__()
        self.fc1 = nn.Linear(lat_size,hid_size)
        self.fc2 = nn.Linear(hid_size,2*hid_size)
        self.fc3 = nn.Linear(2*hid_size,4*hid_size)
        self.fc4 = nn.Linear(4*hid_size,fea_size)
        self.drop = nn.Dropout(p=0.3)

    def forward(self,x):
        x = F.leaky_relu(self.fc1(x),0.2)
        x = self.drop(x)
        x = F.leaky_relu(self.fc2(x),0.2)
        x = self.drop(x)
        x = F.leaky_relu(self.fc3(x),0.2)
        x = self.drop(x)
        x = torch.tanh(self.fc4(x))

        return x


def train(trainLoader,fea_size,hid_size,out_size,lat_size,epoch,lr):
    d_model = Discriminator(fea_size,hid_size,out_size)
    g_model = Generator(lat_size,hid_size,fea_size)
    d_costFunc = nn.BCEWithLogitsLoss()
    g_costFunc = nn.BCEWithLogitsLoss()
    d_optim = optim.Adam(d_model.parameters(),lr)
    g_optim = optim.Adam(g_model.parameters(),lr)

    losses = []
    for e in range(epoch):
        D_loss = 0
        G_loss = 0
        for imags,_ in trainLoader:
            imags = 2*imags - 1
            d_optim.zero_grad()
            d_real_out = d_model.forward(imags)
            d_real_lab = torch.ones(imags.shape[0])
            d_real_loss = d_costFunc(d_real_out.squeeze(),d_real_lab)
            lat_array = np.random.uniform(-1,1,size=(imags.shape[0],lat_size))
            lat_tensor = torch.from_numpy(lat_array).float()
            fake_imgs = g_model.forward(lat_tensor)
            d_fake_out = d_model.forward(fake_imgs)
            d_fake_lab = torch.zeros(imags.shape[0])
            d_fake_loss = d_costFunc(d_fake_out.squeeze(),d_fake_lab)
            d_loss = d_real_loss + d_fake_loss
            d_loss.backward()
            d_optim.step()
            D_loss += d_loss

            g_optim.zero_grad()
            lat_array = np.random.uniform(-1, 1, size=(imags.shape[0], lat_size))
            lat_tensor = torch.from_numpy(lat_array).float()
            fake_imgs = g_model.forward(lat_tensor)
            d_fake_out = d_model.forward(fake_imgs)
            g_loss = g_costFunc(d_fake_out.squeeze(),torch.ones(imags.shape[0]))
            g_loss.backward()
            g_optim.step()
            G_loss += g_loss

        losses.append((D_loss/len(trainLoader),G_loss/len(trainLoader)))
        print("Epoch: {}/{},".format(e,epoch),
              "D_loss: {:.3f},".format(D_loss/len(trainLoader)),
              "G_loss: {:.3f}".format(G_loss/len(trainLoader)))


if __name__ == "__main__":

    lr = 0.001
    epoch = 5
    fea_size = 784
    hid_size = 32
    out_size = 1
    lat_size = 64
    batch_size = 64
    trainLoader,testLoader = get_dataLoader(batch_size)
    train(trainLoader,fea_size,hid_size,out_size,lat_size,epoch,lr)