论文阅读|CycleGAN


CycleGAN论文: https://arxiv.org/pdf/1703.10593.pdf
原作者实现CycleGAN的Pytorch版本: https://github.com/junyanz/pytorch-CycleGAN-and-pix2pix


理解CycleGAN模型和代码

模型优点

CycleGAN是一种基于生成对抗网络(GAN)发展而来的无监督机器学习方法。它是在pix2pix的基础上进行改进和扩展的。CycleGAN主要用于处理非配对图片的图像生成和转换任务,可以实现不同风格之间的转换。

与传统的GAN模型不同,CycleGAN的训练过程中没有配对的数据集。而是通过引入循环一致性损失来实现学习图像之间的映射关系。具体而言,CycleGAN包含两个生成器网络和两个判别器网络。其中,一个生成器负责将图像从源领域转换到目标领域,另一个生成器则相反;两个判别器负责判断生成的图像是否真实。通过交替训练生成器和判别器,CycleGAN能够学习到领域之间的映射规律。

CycleGAN广泛应用于图像风格转换任务。例如,它可以将照片转换为油画风格或将橘子转换为苹果,甚至可以进行马和斑马之间的图像转换。由于CycleGAN不需要成对的数据集,因此在数据准备上更加简单,具有很大的应用前景。

模型结构

CycleGAN由两个生成器和两个判别器构成

G_AtoB() 是风格A向风格B的生成网络
G_BtoA() 是风格B向风格A的生成网络
dis_A() 是判别输入图片是否属于风格A的判别网络
dis_B() 是判别输入图片是否属于风格B的判别网络

  1. G_AtoB()和G_BtoA()的输入为[B, C, W, H],即batchsize, channels, width, height,输出与输入相同;
  2. dis_A()和dis_B()的输入为[B, C, W, H],即batchsize, channels, width, height,输出的维度是[B, 1],网络中经过sigmoid函数输出,最后的取值范围在[0, 1]进行分类。

real_A 是从风格A的真实照片
real_B 是从风格B的真实照片
AtoB = G_AtoB(real_A) 是real_A经过生成网络转换得到的风格B的照片
BtoA = G_BtoA(real_B) 是real_B经过生成网络转换得到的风格A的照片

下面是模型结构图

生成器

生成器是由编码器、转换器和解码器组成的。

编码器

编码器由三层卷积网络构成,假设编码器的输入为[1 3 256 256],经过一层卷积层,变成[1 64 256 256],经过第二层卷积层变成[1 128 128 128],经过第三层卷积层变成[1 256 64 64]。

具体代码实现细节如下

nn.Conv2d(3, 64,7,1,3,padding_mode='reflect'),
            nn.InstanceNorm2d(64),
            nn.ReLU(),
            nn.Conv2d(64,128,3,2,1),
            nn.InstanceNorm2d(128),
            nn.ReLU(),
            nn.Conv2d(128,256,3,2,1),
            nn.InstanceNorm2d(256),
            nn.ReLU(),

nn,Conv2d(输入通道数,输出通道数,卷积核大小, 步长,填充大小(默认填充0),填充模式=‘reflect’)其中采用的是InstanceNorm2d,并没有采用Normalization进行归一化。

Batch Normalization是指batchsize图片中的每一张图片的同一个通道一起进行Normalization操作。而Instance Normalization是指单张图片的单个通道单独进行Noramlization操作。

转换器

本文转换器使用的是残差网络,残差网络目前大部分采用基于梯度的BP算法进行优化,该网络通常将输入信号向前传播,然后通过逆向传输误差值并利用梯度法更新参数。

残差网络除了减弱梯度消失外,还可以理解为这是一种自适应深度,也就是网络可以自己调节层数的深浅,至少可以退化为输入,不会变得更糟糕。可以使网络变得更深,更加的平滑,使深度神经网络的训练成为了可能。

原文中的描述是如果输入的图片大小是128x128就用6个残差块,如果图片大小是256x256就用9个残差块,残差网络的输入个输出大小一致,所以都是编码器的[1 256 64 64]

具体代码实现细节如下

class ResidualBlock(nn.Module):
    def __init__(self):
        super(ResidualBlock, self).__init__()

        block = [
            nn.Conv2d(256,256,3,1,1, padding_mode = 'reflect'),
            nn.InstanceNorm2d(256),
            nn.ReLU(),
            nn.Conv2d(256,256,3,1,1, padding_mode = 'reflect'),
            nn.InstanceNorm2d(256),
        ]

        self.block = nn.Sequential(*block)

    def forward(self, x):
        return x + self.block(x)

文中使用9个残差块

for _ in range(9):
            model += [
                ResidualBlock(),
            ]

解码器

解码器中用到的是反卷积(逆卷积)和卷积层,经过残差结构的tensor为[1 256 64 64],经过第一层反卷积得到[1 128 128 128]经过第二层反卷积层得到[1 64 256 256],再经过卷积层得到[1 3 256 256],得到1张3通道的256x256的图片。

具体代码实现细节如下

nn.ConvTranspose2d(256,128,3,2,1,output_padding=1),
            nn.InstanceNorm2d(128),
            nn.ReLU(),
            nn.ConvTranspose2d(128,64,3,2,1,output_padding=1),
            nn.InstanceNorm2d(64),
            nn.ReLU(),
            nn.Conv2d(64,3,7,1,3,padding_mode='reflect'),
            nn.Tanh()

最后经过Tanh映射到[-1 1]上,输出生成器重建的图像。

判别器

判别器使用的是PatchGAN结构,普通的GAN判别器是将输入直接映射成一个表示输入样本是否为真样本的概率值,而PatchGAN是将输入映射为 N×N 的矩阵 X ,其中 X_ij 表示其中一个patch为真实样本的概率,然后将 X_ij 累加再求均值,即得到判别器的最终输出。

输入图像首先经过卷积核大小为4×4,步长为2,滤波器数量分别为64、128、256的三个卷积层,然后经过卷积核大小为4×4,步长为1,滤波器数量分别为512和1的两个卷积层,得到尺寸大小为30×30×1 的张量,最后经过均值池化层求平均值得到大小为 1×1×1 的张量,再通过调整该张量的维度,得到最终判断输入样本为真实样本的概率值 p。

具体代码实现细节如下

nn.Conv2d(3,64,4,2,1),
            nn.LeakyReLU(0.2),
            nn.Conv2d(64,128,4,2,1),
            nn.InstanceNorm2d(128),
            nn.LeakyReLU(0.2),
            nn.Conv2d(128,256,4,2,1),
            nn.InstanceNorm2d(256),
            nn.LeakyReLU(0.2),
            nn.Conv2d(256,512,4,1,1),
            nn.InstanceNorm2d(512),
            nn.LeakyReLU(0.2),
            nn.Conv2d(512,1,4,1,1)

最后再经过:

F.avg_pool2d(x, x.size()[2:]).view(x.size()[0], -1)

训练细节

数据读取

class ImageDataset(Dataset):
    def __init__(self, root = arg.train_dataroot, unaligned=False):
        self.transform = transforms.Compose([
        transforms.Resize(int(arg.size*1.2), interpolation = Image.BICUBIC), #调整输入图片的大小
        transforms.RandomCrop(arg.size), #随机裁剪
        transforms.RandomHorizontalFlip(),#随机水平翻转图像
        transforms.ToTensor(),
        transforms.Normalize((0.5,0.5,0.5), (0.5,0.5,0.5))  
    ])
        self.unaligned = unaligned

        self.files_A = sorted(glob.glob(root + (arg.train_dataroot.split('/')[0]).split('2')[0] + '/*.jpg'))
        self.files_B = sorted(glob.glob(root + (arg.train_dataroot.split('/')[0]).split('2')[-1] + '/*.jpg'))

    def __getitem__(self, index):
        item_A = self.transform(Image.open(self.files_A[index % len(self.files_A)]).convert('RGB'))

        if self.unaligned:
            item_B = self.transform(Image.open(self.files_B[random.randint(0, len(self.files_B) - 1)]).convert('RGB'))
        else:
            item_B = self.transform(Image.open(self.files_B[index % len(self.files_B)]).convert('RGB'))

        return {'A': item_A, 'B': item_B}

    def __len__(self):
        return max(len(self.files_A), len(self.files_B))

def train_data_loader():
    train_data_loader = DataLoader(ImageDataset(unaligned=True),batch_size=arg.batchSize, shuffle=True, pin_memory=True, drop_last=True)
    return train_data_loader

训练数据集文件的存放要求:

-A2B
----train
----------A
--------------A_images
.
.
.
----------B
--------------B_images
.
.
.

通过调用train_data_loader()函数得到字典格式的数据,可以通过data[‘A’],和data[‘B’]操作将不同类型的图片取出来。

其中的图片还会经过:

调整图片大小至[1.2size 1.2size]
随机裁减至[size size]大小
随机水平反转
归一化

生成图片缓冲区

为了使训练模型更加稳定,实验设置了一个图像缓冲区,用来存储在训练过程中生成器生成的动漫风格图像。在单独训练判别器时,随机混合加载缓冲区里面的动漫风格图像和当前轮生成器最新生成的图像。该缓冲区的大小由变量max_size控制,本实验中默认设置为50。通过传入这种不同时间生成的动漫风格图像给判别器进行训练,可以极大地提升模型训练的稳定性,保证损失函数的值不会大幅度的上下波动。

具体代码实现细节如下

class ReplayBuffer():
    def __init__(self, max_size=50):
        assert (max_size > 0), 'Empty buffer or trying to create a black hole. Be careful.'
        self.max_size = max_size
        self.data = []
        

    def push_and_pop(self, data):
        to_return = []
        for element in data.data:
            element = torch.unsqueeze(element, 0)
            if len(self.data) < self.max_size:
                self.data.append(element)
                to_return.append(element)
            else:
                if random.uniform(0,1) > 0.5:
                    i = random.randint(0, self.max_size-1)
                    data_index.append(i)
                    to_return.append(self.data[i].clone())
                    self.data[i] = element
                else:
                    to_return.append(element)
        return Variable(torch.cat(to_return))

缓冲区的数据初始化大小为50,当缓冲区没有图片的时候,我们把输入的data写入缓冲区,并且返回输入图片,当缓冲区满的时候,50%的可能会随机更新缓冲区数据,将新的数据放进来,替换掉之前生成的数据,之前的数据返回,也会有50%的可能直接返回输入的data数据。

更新学习率

学习率初始为0.0002,总的epoch为200,在0-100的时候,学习率为0.0002,在100-200的时候,学习率逐渐线性减小为0,所以需要进行学习率的更新。

pytorch中提供了torch.optim.lr_scheduler.LambdaLR()函数,其中的学习率衰减需要自己编写函数设定。

利用python实现为:

class MyLambdaLR():
    def __init__(self, n_epochs, offset, decay_start_epoch):
        self.n_epochs = n_epochs
        self.offset = offset
        self.decay_start_epoch = decay_start_epoch

    def step(self, epoch):
        return 1.0 - max(0, epoch + self.offset - self.decay_start_epoch)/(self.n_epochs - self.decay_start_epoch)

模型初始化

在第一次训练的时候对模型中参数进行初始化。学习率的初始值为0.0002,模型会迭代训练200个epoch,因此,我们考虑在第0-100个epoch期间,学习率为0.0002,在训练达到第100个epoch时,学习率将开始线性减小,直到减小为0。

参数初始化代码:

def weights_init_normal(m):
    classname = m.__class__.__name__
    if classname.find('Conv') != -1:
        torch.nn.init.normal_(m.weight.data, 0.0, 0.02)
    elif classname.find('BatchNorm2d') != -1:
        torch.nn.init.normal_(m.weight.data, 1.0, 0.02)
        torch.nn.init.constant_(m.bias.data, 0.0)

保存图片

在训练中得到的结果都是tensor,如何由张量得到图片进行存储和查看,也是十分重要。下面的代码使gpu上的 -1~1 之间的数据转化为0-255之间的值。

具体代码实现细节如下

def TensorToImage(T):
    real_image = 255*(T.cpu().float().numpy()*0.5 + 0.5)
    real_image = real_image.astype(np.uint8).transpose(1, 2, 0)
    real_image = real_image[:,:,[2,1,0]]
    return real_image

文章作者: Wei Hu
文章链接: https://heyhw.cn
版权声明: 本博客所有文章除特別声明外,均采用 CC BY 4.0 许可协议。转载请注明来源 Wei Hu !
  目录