CycleGAN的pytorch实现与概述

作者: MoBaiGeneral 分类: 对抗生成网络 发布时间: 2021-10-18 18:37 访问量:13,564
FavoriteLoading收藏

2021-10-19

本文使用pytorch 1.9.0、cuda 11.1框架,11th Gen Intel(R) Core(TM) i7-11700KF @ 3.60GHz 八核处理器,于window10 64位系统上实现基于horse2zebra数据集的CycleGAN网络搭建与训练(源码可见CycleGAN实现)。想要近一步了解CycleGAN网络,可直接访问网址或查阅论文


Step 1:数据预处理

本文采用horse2zebra数据集(可获取于邮箱mobaisuky@163.com),包括1334个斑马图像数据和1067个马图像数据,斑马图像与马图像并非不同像域内的成对映射,图像分辨率为256*256,图像格式为jpg。在数据预处理阶段,将图像读取为具三通道的三维数据,通过torchvision.transforms工具对图像数据进行一系列增强处理,如数据归一化等。


Step 2:模型定义

网络模型主要包括生成器和判别器,而生成器由下采样、残差网络、上采样三个部分构成;差别器则由一系列的卷积、归一化、激活层三层组合构成,pytorch的归一化层主要有BatchNorm、LayerNorm、InstanceNorm、GroupNorm等。其中,BatchNorm系batch方向做归一化,算NHW的均值,对小batchsize效果不好,其主要缺点是对batchsize的大小比较敏感,由于每次计算均值和方差是在一个batch上,所以如果batchsize太小,则计算的均值、方差不足以代表整个数据分布;LayerNorm系channel方向做归一化,算CHW的均值,主要对RNN作用明显;InstanceNorm系一个channel内做归一化,算H*W的均值,用在风格化迁移。因为在图像风格化中,生成结果主要依赖于某个图像实例,所以对整个batch归一化不适合图像风格化中,因而对HW做归一化。可以加速模型收敛,并且保持每个图像实例之间的独立。GroupNorm将channel方向分group,然后每个group内做归一化,算(C//G)HW的均值;这样与batchsize无关,不受其约束;SwitchableNorm是将BN、LN、IN结合,赋予权重,让网络自己去学习归一化层应该使用什么方法。本程序中归一化方式采用的是通道内归一化。


Step 3:训练框架

训练模型的基本步骤主要由optimizer.zero_grad()、loss.backward()、optimizer.step()、lr_scheduler.step()依次组成。其中,函数optimizer.zero_grad()将梯度归零,然后函数loss.backward()反向传播计算得到每个参数的梯度值,最后通过函数optimizer.step()实现梯度下降以执行一步参数更新。训练的过程通常使用mini-batch方法,如果不将梯度清零,梯度会与上一个batch的数据相关,则optimizer.zero_grad()函数要写在反向传播和梯度下降之前。损失函数loss是由模型的所有权重w经过一系列运算得到,若某个w的requires_grads为True,则w的所有上层参数(后面层的权重w)的grad_fn属性中将保存对应的运算,在使用loss.backward()后会一层层的反向传播计算每个w的梯度值,并保存到该w的grad属性中。如果没有进行tensor.backward(),则梯度值将会是None,因此loss.backward()要先于optimizer.step()之前。step()函数的作用是执行一次优化步骤,通过梯度下降法来更新参数的值。因为梯度下降是基于梯度的,所以在执行optimizer.step()函数前应先执行loss.backward()函数来计算梯度。值得注意的是,optimizer只负责通过梯度下降进行优化,而不负责产生梯度,梯度由tensor.backward()方法产生。torch.optim.lr_scheduler模块提供了一些根据epoch训练次数来调整学习率(learning rate)的方法。一般情况下我们会设置随着epoch的增大而逐渐减小学习率从而达到更好的训练效果,学习率的调整应该放在optimizer更新之后(注意: 在PyTorch 1.1.0之前的版本,学习率的调整应该被放在optimizer更新之前的)。optimizer具有属性optimizer.defaults和optimizer.param_groups。其中,optimizer.defaults是字典,存放这个优化器的一些初始参数,有:’lr’, ‘betas’, ‘eps’, ‘weight_decay’, ‘amsgrad’。而optimizer.param_groups则是列表,每个元素都是一个字典,每个元素包含的关键字有:’params’, ‘lr’, ‘betas’, ‘eps’, ‘weight_decay’, ‘amsgrad’,params类是各个网络的参数放在了一起。lr_scheduler更新optimizer的lr,即更新的optimizer.param_groups[n][‘lr’],而不是optimizer.defaults[‘lr’]。

训练框架伪代码如下所示:

定义数据集(可继承Dataset)以实现数据加载器DataLoader;

定义net并初始化权重;

定义optimizer,lr_scheduler以及损失计算方法criterion;

for epoch=startEpoch:endEpoch:

        for batch=0:nBatch:

                optimizer.zero_grad()

                output=net(input);

                loss=criterion(output,real)

                loss.backward()

                optimizer.step()

                lr_schedule.step()

模型的训练过程旨在优化两项目标损失:

其一,对抗损失(adversarial loss):促使生成图像分布与目标域的图像分布相逼近。其二,循环一致性损失(cycle consistency loss):使得F(G(x))≈x以及G(F(x))≈y,其中:G:XYF:YX,能够防止学习到的映射G和F互相矛盾。


Step 4:效果

在程序中,将模型损失细划分为由两个生成器损失,两个针对生成结果的判别损失,两个循环一致性损失的加权组合构成的损失G,判别器A针对真实结果和异构生成结果的判别损失D_A和判别器B针对真实结果和异构生成结果的判别损失D_B三个部分,并定义三倍数量的optimizer.和lr_scheduler对这三项损失进行优化训练。

通过配置参数批次大小2,图像分辨率256*256,初始学习率0.0008,训练轮次128对模型进行训练学习。在学习过程中,三项损失的变化趋势如下所示。在训练的前10轮内,模型的各项损失均有指数形式的下降趋势,后则逐步趋于平缓。然则,在迭代学习80轮次后,损失G下降变得较为困难,判别损失的下降幅度依旧可观,表明生成器的学习效果逐渐较判别器拉开了距离,有待进一步的优化处理。

     

对训练得到的模型效果进行测试。

选取斑马图像,输入训练得到的生成器G_A对马进行生成,其效果如下所示: 

 

 

 

选取马图像,输入训练得到的生成器G_B对斑马进行生成,其效果如下所示:

 

 


参考文献:

https://blog.csdn.net/shanglianlm/article/details/85075706

https://www.cnblogs.com/Thinker-pcw/p/9630367.html

https://blog.csdn.net/cassiepython/article/details/80942899

     

如果觉得小墨的文章对您有用,请随意打赏。您的支持将鼓励我继续创作!

一条评论

发表评论