生成器
我们的生成器模型将完全遵循DCGAN的完全相同的体系结构。
与生成器类似，但是相反，我们不是将图像卷积为一组特征，而是采用随机噪声输入“ z”并运行一系列反卷积以达到所需的图像形状（3，64，64 ）。
classGenerator(nn.Module): def__init__(self, channels_noise, channels_img, features_g): super(Generator, self).__init__() self.net=nn.Sequential( #Input: Nxchannels_noisex1x1self._block(channels_noise, features_g*16, 4, 1, 0), #img: 4x4self._block(features_g*16, features_g*8, 4, 2, 1), #img: 8x8self._block(features_g*8, features_g*4, 4, 2, 1), #img: 16x16self._block(features_g*4, features_g*2, 4, 2, 1), #img: 32x32nn.ConvTranspose2d( features_g*2, channels_img, kernel_size=4, stride=2, padding=1 ), #Output: Nxchannels_imgx64x64nn.Tanh(), ) def_block(self, in_channels, out_channels, kernel_size, stride, padding): returnnn.Sequential( nn.ConvTranspose2d( in_channels, out_channels, kernel_size, stride, padding, bias=False, ), nn.BatchNorm2d(out_channels), nn.ReLU(), ) defforward(self, x): returnself.net(x)
最后，我们只需要在训练之前将initialize_weights（）函数添加到两个模型即可：
definitialize_weights(model): #InitializesweightsaccordingtotheDCGANpaperforminmodel.modules(): ifisinstance(m, (nn.Conv2d, nn.ConvTranspose2d, nn.BatchNorm2d)): nn.init.normal_(m.weight.data, 0.0, 0.02)
训练
现在我们有了数据和模型，我们可以从train.py文件开始，在这里我们将加载数据，初始化模型并运行主训练循环。
超参数和导入
importnumpyasnpimporttorchimporttorch.nnasnnimporttorch.optimasoptimimporttorchvisionimporttorchvision.datasetsasdatasetsimporttorchvision.transformsastransformsfromtorch.utils.dataimportDatasetfromtorch.utils.dataimportDataLoaderfromtorch.utils.tensorboardimportSummaryWriterfrommodelimportDiscriminator, Generator, initialize_weightsimportrandomimportosimportnatsortfromPILimportImage, ImageOps, ImageEnhancedevice=torch.device("cuda"iftorch.cuda.is_available() else"cpu") D_LEARNING_RATE=2e-4G_LEARNING_RATE=1e-4BATCH_SIZE=64IMAGE_SIZE=64CHANNELS_IMG=3NOISE_DIM=128NUM_EPOCHS=100FEATURES_DISC=64FEATURES_GEN=64
两个优化器可以使用相同的学习率，但是我发现对鉴别器使用稍高的学习率被证明是更有效的。在更复杂的数据集上，我发现较小的批次大小（例如16或8）可以帮助避免过度拟合。
随机增强
改善GAN训练并从数据集中获得最大收益的技术之一是应用随机图像增强。在原始论文中，它们还提供了一种在生成器端还原增强图像的机制，因为我们不希望生成器生成增强图像。但是，在这种情况下，我认为应用这些简单的增幅就足够了，而这些增幅并不会真正影响画质。如果我们试图获得照片般逼真的结果，那么使用它来进行全面实施可能是一个更好的主意。
defrandom_augmentation(img): #randommirroring/flippingimagerand_mirror=random.randint(0,1) #randomsaturationadjustmentrand_sat=random.uniform(0.5,1.5) #randomsharpnessadjustmentrand_sharp=random.uniform(0.5,1.5) converter=ImageEnhance.Color(img) img=converter.enhance(rand_sat) converter=ImageEnhance.Sharpness(img) img=converter.enhance(rand_sharp) ifrand_mirror==0: img=ImageOps.mirror(img) returnimg
我们对图像执行3个操作-镜像，饱和度调整和锐度调整。镜像图像对我们的图像质量没有影响，因为我们只是在翻转图像。对于饱和度和清晰度，我使用了一个较小的系数范围（0.5、1.5），以免对原始图像造成很大的影响。
数据加载器
为了应用我们之前构建的随机增强方法并加载数据，我编写了一个使用其下定义的转换的自定义数据集。
classCustomDataSet(Dataset): def__init__(self, main_dir, transform): self.main_dir=main_dirself.transform=transformall_imgs=os.listdir(main_dir) self.total_imgs=natsort.natsorted(all_imgs) def__len__(self): returnlen(self.total_imgs) def__getitem__(self, idx): img_loc=os.path.join(self.main_dir, self.total_imgs[idx]) image=Image.open(img_loc).convert('RGB') image=random_augmentation(image) tensor_image=self.transform(image) returntensor_imagetransforms=transforms.Compose( [ transforms.Resize((IMAGE_SIZE,IMAGE_SIZE)), transforms.ToTensor(), transforms.Normalize( [0.5for_inrange(CHANNELS_IMG)], [0.5for_inrange(CHANNELS_IMG)] ), ] ) dataset=CustomDataSet("./data/128_portraits/", transform=transforms) dataloader=DataLoader(dataset, batch_size=BATCH_SIZE, shuffle=True)
训练
最后，我们可以初始化我们的网络并开始对其进行训练。对于鉴别器训练，我使用均方误差作为损失函数。我也尝试使用二进制交叉熵，但MSELoss最有效。在训练循环之前，我们还初始化张量板编写器以在tensorboard上实时查看我们的图像。
gen=Generator(NOISE_DIM, CHANNELS_IMG, FEATURES_GEN).to(device) disc=Discriminator(CHANNELS_IMG, FEATURES_DISC).to(device) initialize_weights(gen) initialize_weights(disc) ###uncommenttoworkfromsavedmodels####gen.load_state_dict(torch.load('saved_models/generator_model.pt')) #disc.load_state_dict(torch.load('saved_models/discriminator_model.pt')) opt_gen=optim.Adam(gen.parameters(), lr=G_LEARNING_RATE, betas=(0.5, 0.99)) opt_disc=optim.Adam(disc.parameters(), lr=D_LEARNING_RATE, betas=(0.5, 0.99)) criterion=nn.MSELoss() fixed_noise=torch.randn(16, NOISE_DIM, 1, 1).to(device) writer_real=SummaryWriter(f"logs/real") writer_fake=SummaryWriter(f"logs/fake") step=0gen.train() disc.train() forepochinrange(NUM_EPOCHS): forbatch_idx, realinenumerate(dataloader): real=real.to(device) noise=torch.randn(BATCH_SIZE, NOISE_DIM, 1, 1).to(device) fake=gen(noise) ###TrainDiscriminatordisc_real=disc(real).reshape(-1) loss_disc_real=criterion(disc_real, torch.ones_like(disc_real)) disc_fake=disc(fake.detach()).reshape(-1) loss_disc_fake=criterion(disc_fake, torch.zeros_like(disc_fake)) loss_disc= (loss_disc_real+loss_disc_fake) /2disc.zero_grad() loss_disc.backward() opt_disc.step() ###TrainGeneratorusingfeaturematchingoutput=disc(fake).reshape(-1) loss_gen=criterion(output, torch.ones_like(output)) gen.zero_grad() loss_gen.backward() opt_gen.step() #Printlossesoccasionallyandprinttotensorboardifbatch_idx%10==0: torch.save(gen.state_dict(), 'generator_model.pt') torch.save(disc.state_dict(), 'discriminator_model.pt') print( f"Epoch [{epoch}/{NUM_EPOCHS}] Batch {batch_idx}/{len(dataloader)} \Loss D: {loss_disc:.4f}, loss G: {loss_gen:.4f}" ) withtorch.no_grad(): fake=gen(fixed_noise) img_grid_real=torchvision.utils.make_grid( real[:16], normalize=True ) img_grid_fake=torchvision.utils.make_grid( fake[:16], normalize=True ) writer_real.add_image("Real", img_grid_real, global_step=step) writer_fake.add_image("Fake", img_grid_fake, global_step=step) step+=1
在训练的第一部分中，我们使用MSELoss在真实图像和伪图像上训练鉴别器。之后，我们使用特征匹配来训练我们的生成器。之前，我们在鉴别器的前向传递中添加了变量“ feature_matching”，以从图像中提取感知特征。在传统的DCGAN中，您只需训练生成器以伪造的图像来欺骗鉴别器，而在这里，我们试图训练生成器以生成与真实图像的特征紧密匹配的图像。此技术通常可以提高训练的稳定性。
经过100个批次后，我获得了以下结果。我尝试对模型进行更多的迭代训练，但是图像质量没有太大改善。
结论与最终想法
本文的目的是记录我从事该项目的过程。尽管在线上有很多资源和论文探讨了这个令人兴奋的概念的不同方面，但我发现有些东西是只能通过经验学习的……与其他任何东西一样。但我希望您能在本文中找到一些可以在自己的GAN项目中应用或试验的东西。由于我们获得的结果并不完美，因此我打算应用本文中提出的EvolGAN来优化我的生成器。
