1. 简介
图像风格迁移是计算机视觉领域中的一个热门研究方向,其中 CycleGAN (循环对抗生成网络) 在无监督领域取得了显著的突破。与传统需要成对训练数据的模型如 Pix2Pix 不同,CycleGAN 不需要严格的成对数据,只需两类图片域数据,便可实现图像风格的迁移与互换。
本篇博文将通过一个实际案例演示如何使用 CycleGAN 实现从草图到目标线稿图的图像风格迁移任务,并详细介绍 CycleGAN 的模型结构、数据处理及训练过程。
2. 模型介绍
CycleGAN 的核心思想源自 "Unpaired Image-to-Image Translation using Cycle-Consistent Adversarial Networks" 论文。该模型在不需要成对示例的情况下,学习将源域 X 的图像转换到目标域 Y。其应用领域包括风格迁移、图像增强和域适应等任务。
2.1 CycleGAN 网络结构
CycleGAN 由两个 GAN 模型组成,其对称的架构允许在不同的域之间来回转换图像。具体而言,CycleGAN 使用两个生成器(G 和 F)和两个判别器(D_X 和 D_Y),生成器负责将域 X 的图像转换到域 Y,并通过判别器对生成结果进行真假判断。
模型架构如下:
- 生成器:生成器采用 ResNet 结构,由 9 个残差块组成,适合处理 256x256 尺寸的图片。
- 判别器:判别器通过 PatchGAN 模型检测图像的真实性,以保证生成的图像足够逼真。
2.2 循环一致性损失
CycleGAN 通过 循环一致性损失 来保证从域 X 到域 Y,再从域 Y 转换回域 X 的图像应尽可能接近原始图像。这种损失机制确保模型不会丢失重要的图像特征。
3. 数据集
本案例使用的数据集包含线稿图和草图图像,所有图片大小为 256x256 像素。数据集分为训练集和测试集,训练集包含 25654 张图片,测试集包含约 100 张线稿图片和 116 张草图图片。
4. 模型实现
4.1 生成器模型
生成器模型基于 ResNet 结构,通过卷积、反卷积及残差块实现图像风格的转换。以下是生成器的代码实现:
import mindspore.nn as nn
class ResidualBlock(nn.Cell):
def __init__(self, dim):
super(ResidualBlock, self).__init__()
self.conv_block = nn.SequentialCell(
nn.Conv2d(dim, dim, kernel_size=3, padding=1, pad_mode="pad"),
nn.BatchNorm2d(dim),
nn.ReLU(),
nn.Conv2d(dim, dim, kernel_size=3, padding=1, pad_mode="pad"),
nn.BatchNorm2d(dim)
)
def construct(self, x):
return x + self.conv_block(x)
class ResNetGenerator(nn.Cell):
def __init__(self, input_nc, output_nc, n_residual_blocks=9):
super(ResNetGenerator, self).__init__()
model = [
nn.Conv2d(input_nc, 64, kernel_size=7, padding=3, pad_mode="pad"),
nn.BatchNorm2d(64),
nn.ReLU()
]
# Downsampling
model += [
nn.Conv2d(64, 128, kernel_size=3, stride=2, padding=1),
nn.BatchNorm2d(128),
nn.ReLU(),
nn.Conv2d(128, 256, kernel_size=3, stride=2, padding=1),
nn.BatchNorm2d(256),
nn.ReLU()
]
# Residual blocks
for _ in range(n_residual_blocks):
model += [ResidualBlock(256)]
# Upsampling
model += [
nn.Conv2dTranspose(256, 128, kernel_size=3, stride=2, padding=1, output_padding=1),
nn.BatchNorm2d(128),
nn.ReLU(),
nn.Conv2dTranspose(128, 64, kernel_size=3, stride=2, padding=1, output_padding=1),
nn.BatchNorm2d(64),
nn.ReLU()
]
model += [
nn.Conv2d(64, output_nc, kernel_size=7, padding=3, pad_mode="pad"),
nn.Tanh()
]
self.model = nn.SequentialCell(model)
def construct(self, x):
return self.model(x)
4.2 判别器模型
判别器基于 PatchGAN 的结构,通过卷积网络将输入图片划分为多个小的 patch,并分别进行真假判别。
class Discriminator(nn.Cell):
def __init__(self, input_nc, ndf=64):
super(Discriminator, self).__init__()
self.model = nn.SequentialCell([
nn.Conv2d(input_nc, ndf, kernel_size=4, stride=2, padding=1),
nn.LeakyReLU(0.2),
nn.Conv2d(ndf, ndf * 2, kernel_size=4, stride=2, padding=1),
nn.BatchNorm2d(ndf * 2),
nn.LeakyReLU(0.2),
nn.Conv2d(ndf * 2, ndf * 4, kernel_size=4, stride=2, padding=1),
nn.BatchNorm2d(ndf * 4),
nn.LeakyReLU(0.2),
nn.Conv2d(ndf * 4, ndf * 8, kernel_size=4, stride=1, padding=1),
nn.BatchNorm2d(ndf * 8),
nn.LeakyReLU(0.2),
nn.Conv2d(ndf * 8, 1, kernel_size=4, stride=1, padding=1)
])
def construct(self, x):
return self.model(x)
4.3 优化器与损失函数
CycleGAN 采用对抗性损失和循环一致性损失的组合来训练生成器和判别器。优化器选择了 Adam 优化器,学习率设置为 0.0002。
import mindspore as ms
# 定义损失函数和优化器
gan_loss = nn.BCELoss()
cycle_loss = nn.L1Loss()
optimizer_G = nn.Adam(generator.parameters(), learning_rate=0.0002)
optimizer_D = nn.Adam(discriminator.parameters(), learning_rate=0.0002)
5. 训练与推理
训练过程中,我们交替训练生成器和判别器。判别器通过真假样本的判别进行训练,而生成器则通过对抗判别和循环一致性进行优化。以下是一个训练步骤的实现:
def train_step(real_A, real_B):
# 生成器前向计算
fake_B = generator_A2B(real_A)
fake_A = generator_B2A(real_B)
# 判别器前向计算
D_A_loss = gan_loss(discriminator_A(fake_A), Tensor(0)) + gan_loss(discriminator_A(real_A), Tensor(1))
D_B_loss = gan_loss(discriminator_B(fake_B), Tensor(0)) + gan_loss(discriminator_B(real_B), Tensor(1))
# 生成器损失计算
cycle_A_loss = cycle_loss(generator_B2A(fake_B), real_A)
cycle_B_loss = cycle_loss(generator_A2B(fake_A), real_B)
G_loss = cycle_A_loss + cycle_B_loss + D_A_loss + D_B_loss
optimizer_G.step()
optimizer_D.step()
return G_loss, D_A_loss, D_B_loss
结语
通过本次的CycleGAN模型实践,我们深入理解了图像风格迁移的基本原理,特别是在无监督情况下如何实现两个域之间的图像转换。CycleGAN的循环一致性损失在保持图像内容一致性的同时,又能实现风格的转换,这是其在域迁移任务中广泛应用的重要原因。在整个实现过程中,不仅对生成器和判别器的构建有了更清晰的理解,同时也进一步熟悉了损失函数的优化策略。
这次实验的关键在于让模型具备在没有配对数据的情况下,也能够进行风格转换的能力。虽然实验需要较大的计算资源,但我们通过小规模数据集也能够体验到CycleGAN的强大之处。希望通过这个项目,我们不仅能掌握CycleGAN的基本原理,也能为以后的图像生成和风格迁移任务打下坚实的基础。
如果你觉得这篇博文对你有帮助,请点赞、收藏、关注我,并且可以打赏支持我!
欢迎关注我的后续博文,我将分享更多关于人工智能、自然语言处理和计算机视觉的精彩内容。
谢谢大家的支持!