GAN源码的讲解

核心思想

构建生成器与判别器，生成器用于将随机的噪声样本转换为我们希望的真实样本，判别器用于给伪造的生成样本进行打分，以达到生成随机的真实样本，增加模型样本容量，这常常用于特定目标检测的任务当中，属于一种分类任务。

生成器

# 生成器，用于将随机的噪声样本转换成一张图像
class Generator(nn.Module):
    def __init__(self,img_shape):
        super(Generator, self).__init__()

        def block(in_feat, out_feat, normalize=True):
            layers = [nn.Linear(in_feat, out_feat)]
            if normalize:
                layers.append(nn.BatchNorm1d(out_feat, momentum=0.8))
            layers.append(nn.LeakyReLU(0.2, inplace=True))
            return layers

        self.model = nn.Sequential(
            *block(opt.latent_dim, 128, normalize=False),
            *block(128, 256,False),
            *block(256, 512,False),
            *block(512, 1024,False),
            nn.Linear(1024, int(np.prod(img_shape))),
            nn.Tanh()
        )

    def forward(self, z):
        img = self.model(z)
        img = img.view(img.size(0), *img_shape)
        return img

判别器

# 判别器，用于评判生成样本与目标样本之间的相似度，（会给出具体的得分大小）
class Discriminator(nn.Module):
    def __init__(self,img_shape):
        super(Discriminator, self).__init__()

        self.model = nn.Sequential(
            nn.Linear(int(np.prod(img_shape)), 512),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Linear(512, 256),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Linear(256, 1),
            nn.Sigmoid(),
        )

    def forward(self, img):
        img_flat = img.reshape(img.size(0), -1)
        validity = self.model(img_flat)

        return validity

训练过程

构建二分类损失函数、生成器、判别器

adversarial_loss = torch.nn.BCELoss()
generator = Generator()
discriminator = Discriminator()

加载数据集

os.makedirs("../../data/mnist", exist_ok=True)
dataloader = torch.utils.data.DataLoader(
    datasets.MNIST(
        "../../data/mnist",
        train=True,
        download=True,
        transform=transforms.Compose(
            [transforms.Resize(opt.img_size), transforms.ToTensor(), transforms.Normalize([0.5], [0.5])]
        ),
    ),
    batch_size=opt.batch_size,
    shuffle=True,
)

构建生成器与判别器的优化器

optimizer_G = torch.optim.Adam(generator.parameters(), lr=opt.lr, betas=(opt.b1, opt.b2))
optimizer_D = torch.optim.Adam(discriminator.parameters(), lr=opt.lr, betas=(opt.b1, opt.b2))

训练

for i in range(500):
    #训练生成器
    optimizer_G.zero_grad()#梯度置零
	#随机的噪声样本
	 z = Variable(Tensor(np.random.normal(0, 1, (imgs.shape[0], opt.latent_dim))))
    # 生成一批图片
    gen_imgs = generator(z)
    '''
    #将生成器得到的生成样本送入到判别器当中，得到一批图像对应的得分
    #再将得到的得分与我们认为它是真实样本的分数标签值计算得分损失
    #训练生成器能够更好的产生伪造样本
    #需要注意的是在此期间判别器的参数是不变的，也就是说它此时的判别规则或者说是判别能力没有变
    '''
    g_loss = adversarial_loss(discriminator(gen_imgs), valid)
    g_loss.backward()
    optimizer_G.step()

    #训练判别器
    optimizer_D.zero_grad()

    # Measure discriminator's ability to classify real from generated samples
    '''
    #将真实样本送入到判别器当中计算得分
    #再将得分大小与认为它是满分情况计算损失大小
    #以及再将生成器产生的伪造样本送入到判别器当中计算得分，与判别器认为它是假样本计算损失
    #训练判别器提高真实样本的判别得分外，还要提高其识别假样本的能力
    #同样需要注意的是在此期间生成器的阐述是不变的，也就是此时生成器的伪造能力并没有得到提高
    '''
    real_loss = adversarial_loss(discriminator(real_imgs), valid)
    fake_loss = adversarial_loss(discriminator(gen_imgs.detach()), fake)
    d_loss = (real_loss + fake_loss) / 2
    d_loss.backward()
    optimizer_D.step()
    #保存生成的样本
    if i%100==0:
     save_image(gen_imgs.data, "images/%d.png" % batches_done, nrow=5, normalize=True)