用预训练的卷积神经网络完成色彩风格迁移

用预训练的卷积神经网络完成色彩风格迁移

摄影爱好者也许接触过滤波器。它能改变照片的颜色风格, 从而使风景照更加锐利或者令人像更加美白。 但一个滤波器通常只能改变照片的某个方面。如果要照片达到理想中的风格, 可能需要尝试大量不同的组合。这个过程的复杂程度不亚于模型调参。

输入内容图像和风格图像,输出风格迁移后的合成图像

卷积神经网络简介

卷积神经网络 (Convolutional Neural Network, CNN) 是一种深度学习模型,特别适用于处理具有网格结构的数据,如图像和视频。 CNN 主要由卷积层 (Convolutional Layer)、池化层 (Pooling Layer)、激活函数 (Activation Function) 和全连接层 (Fully Connected Layer) 组成。

在 CNN 中,卷积层通过卷积操作提取输入数据的特征,池化层通过降采样操作减小特征图的大小,激活函数引入非线性因素,全连接层将提取的特征映射到输出类别。 整个网络通过反向传播算法进行训练,优化网络参数以最小化损失函数。

CNN 的特点包括参数共享、局部感受野和空间层次结构,使其在图像处理任务中表现出色。 通过多层卷积和池化操作,CNN 能够逐渐学习到输入数据的高级特征表示,从而实现图像分类、目标检测、图像分割等任务。

CNN 的数学原理涉及卷积操作、池化操作、激活函数等,其中卷积操作是 CNN 的核心部分。卷积操作可以用数学公式表示为: \[ \text{Conv}(I, K) = I * K \] 其中,\(I\) 表示输入特征图,\(K\) 表示卷积核 (滤波器),\(*\) 表示卷积操作。卷积操作通过滑动卷积核在输入特征图上提取特征,实现特征的局部连接和参数共享。

总的来说,CNN 是一种强大的深度学习模型,在图像处理领域有着广泛的应用,并且在许多其他领域也取得了成功。

VGG19 简介

VGG19 是一种深度卷积神经网络模型,属于 VGG (Visual Geometry Group) 系列之一。 该模型由牛津大学的研究团队开发,用于图像分类和识别任务。VGG19 的名称中的 “19” 表示该模型具有 19 层神经网络,包括 16 个卷积层和 3 个全连接层。

VGG19 的架构相对简单而经典,采用了连续多个 \(3 \times 3\) 的卷积核和 \(2 \times 2\) 的最大池化层,以及 Relu 激活函数。 这种深层的架构使得 VGG19 能够学习到更加复杂和抽象的特征,从而在图像分类等任务中表现出色。

VGG19 模型在图像识别领域取得了很好的表现,并被广泛用于迁移学习和特征提取。它的预训练模型通常基于大规模图像数据集 (如 ImageNet) 进行训练, 可以在各种计算机视觉任务中发挥作用。

代码实现:选择训练设备并设置 VGG 中参数

使用 PyTorch 中的 models.vgg19 函数加载预训练的 VGG19 模型,并将其特征提取部分 features 移动到指定的设备 device 上进行计算, 并将 vgg 模型中的参数的梯度设置为不需要计算,即不进行梯度更新。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
from PIL import Image
import numpy as np
import matplotlib.pyplot as plt

import torch
import torch.optim as optim
from torchvision import models, transforms
from torchvision.models import VGG19_Weights

from tqdm import tqdm

# 选择设备
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
# 使用 PyTorch 中的 models.vgg19 函数加载预训练的 VGG19 模型,并将其特征提取部分 features 移动到指定的设备 device 上进行计算
vgg = models.vgg19(weights=VGG19_Weights.IMAGENET1K_V1).features.to(device)

# 将 vgg 模型中的参数的梯度设置为不需要计算,即不进行梯度更新
for i in vgg.parameters():
i.requires_grad_(False)

载入内容图片与风格图片

在使用 torchvision.transforms 进行数据处理时我们经常进行的操作是:

1
2
3
4
5
6
# 定义读入与处理图片的函数

MIX_SIZE = 500 # 设定图像最大尺寸

rgb_mean = (0.485, 0.456, 0.406)
rgb_std = (0.229, 0.224, 0.225)

第 5 行代码 (0.485,0.456,0.406) 表示均值,分别对应的是 RGB 三个通道;第 6 行代码 (0.229,0.224,0.225) 则表示的是标准差。 均值和标准差的值由 ImageNet 数据集计算得出。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
def load_img(path, max_size=MIX_SIZE, shape=None):
img = Image.open(path).convert("RGB")

if max(img.size) > max_size:
size = max_size
else:
size = max(img.size)

if shape is not None:
size = shape

transform = transforms.Compose(
[
transforms.Resize(size),
transforms.ToTensor(),
transforms.Normalize(mean=rgb_mean, std=rgb_std)
]
)
# 删除 alpha 通道 jpg,转为 png,并补足另一维度 batch
img = transform(img)[:3, :, :].unsqueeze(0)
return img.to(device)
# 载入图像
content = load_img("./style_transfer/IMG_20230414_183148.jpg")
style = load_img(
"./style_transfer/style.jpg", shape=content.shape[-2:]
) # 让两张图尺寸一样

代码实现:将 PyTorch 张量转换为表示图像的 NumPy 数组的函数

将 PyTorch 张量移动到 CPU 上,然后转换为 NumPy 数组,并通过 squeeze() 方法去除维度为 1 的维度,以简化数组的形状; 然后,转置数组的维度,将通道维度移动到最后一个维度,以匹配图像的通道顺序,并对图像进行归一化处理,通过乘以标准差并加上均值来反向转换图像数据; 最后,将图像像素值截断到范围 [0, 1],确保图像数据在有效范围内。

1
2
3
4
5
6
7
8
9
10
11
12
13
# 将 PyTorch 张量转换为表示图像的 NumPy 数组的函数
def im_convert(tensor):
# 创建输入张量的副本,并且将其从计算图中分离,确保不会影响原始张量的梯度计算
img = tensor.clone().detach()
# 将 PyTorch 张量移动到 CPU 上,然后转换为 NumPy 数组,并通过 squeeze() 方法去除维度为 1 的维度,以简化数组的形状
img = img.cpu().numpy().squeeze()
# 转置数组的维度,将通道维度移动到最后一个维度,以匹配图像的通道顺序
img = img.transpose(1, 2, 0)
# 对图像进行归一化处理,通过乘以标准差并加上均值来反向转换图像数据
img = img * np.array(rgb_std) + np.array(rgb_mean)
# 将图像像素值截断到范围 [0, 1],确保图像数据在有效范围内
img = img.clip(0, 1)
return img

代码实现:将网络层层的输出存入 feature 数组

这段代码是一个函数,名为 get_features,它用于将给定图像在神经网络模型中经过指定层的输出存储到一个名为 features 的字典中,并返回该字典。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
# 将网络层层的输出存入 feature 数组中,并返回
def get_features(img, model, layers=None):
if layers is None:
layers = {
"0": "conv1_1",
"5": "conv2_1",
"10": "conv3_1",
"19": "conv4_1",
"21": "conv4_2", # content 层所需
"28": "conv5_1",
}

features = {}
x = img
for name, layer in model._modules.items():
x = layer(x)
if name in layers:
features[layers[name]] = x

return features

函数接受三个参数: img 代表输入的图像数据,model 代表神经网络模型,layers 是一个字典,用于指定要提取特征的层。 如果未指定 layers,则会使用默认的层次对应关系。

函数首先创建一个空字典 features 用于存储特征。然后,对输入图像 img 进行前向传播计算,通过遍历模型的各个层, 将每一层的输出存储到 features 字典中,键为层的名称,值为该层的输出。

最后,函数返回存储有指定层输出的 features 字典。

Gram 矩阵

设提取内容的模型为 \(f_c\),输入图像的矩阵为 \(\vb*{X}\),内容图像的矩阵为 \(\vb*{C}\),我们直接用平方误差作为内容上的损失: \[ L_c(\vb*{X})=\dfrac{1}{2}\norm{f_c(\vb*{X})-f_c(\vb*{C})}^2_F \] 相比于内容,图像的风格更难描述,也更不容易直接得到。直观上说,一幅图像的风格值的是图像的色彩、纹理等要素,这些要素在每个尺度上都有体现。 因此,我们考虑同一个卷积层中由不同的卷积核提取的特征。由于这些特征属于同一层,因此基本属于原图中的相同尺度。 那么它们之间的相关性可以一定程度上反映图像在该尺度上的风格特点。

设某一层中卷积核是数量为 \(N\),每个卷积核与输入运算得到的输出大小为 \(M\),其中 \(M\) 表示输出矩阵中元素的总数量。 那么,该卷积层中的总特征矩阵 \(\vb*{F}\in\mathbb{R}^{N\times M}\),每一行向量都表示一个卷积核输出的特征。

为了表示不同特征之间的关系,我们使用与因子分解机类似的思想,将特征之间做内积,得到 \[ \vb*{G}=\vb*{FF}^\top \] 该矩阵为格拉姆矩阵 (Gram matrix),矩阵 \(\vb*{G}\in\mathbb{R}^{N\times N}\) 在计算乘积的过程中,把与特征在图像中的位置有关的维度消掉了,只保留了与卷积核有关的维度。 直观上说,这体现了图像的风格特征与其相对位置无关,是图像的全局属性。

设由第 \(i\) 个卷积层提取的输入图像的风格矩阵为 \(\vb*{G}_X^{(x)}\),风格图像的风格矩阵为 \(\vb*{G}_S^{(i)}\),我们同样用平方误差作为损失函数: \[ L_S^{(i)}(\vb*{X})=\dfrac{1}{4N_{(i)}^2M_{(i)}^2}\norm{\vb*{G}_{X}^{(i)}-\vb*{G}_{S}^{(i)}}^2_F \] 这里,由于不同卷积层的参数可能不同,我们额外除以 \(4N_{(i)}^2M_{(i)}^2\) 来进行层间的归一化。

最后,再用权重 \(w_i\) 对不同卷积层的损失做加权平均,得到总的风格损失为: \[ L_S(\vb*{X})=\sum_{i}w_iL_S^{(i)}(\vb*{X}) \]

1
2
3
4
5
6
7
8
9
10
# 输出 Gram Matrix
def gram_matrix(tensor):
# d 表示通道数,h 表示高度,w 表示宽度。
_, d, h, w = tensor.size() # 第一个是 batch_size

# 将张量重新调整形状为 (d, h*w),即将每个通道的特征图展平为一维向量
tensor = tensor.view(d, h * w)
gram = torch.mm(tensor, tensor.t())

return gram

内容特征与风格特征提取

首先通过 get_features 函数分别获取内容图像 content 和风格图像 style 在 VGG 模型中各层的特征表示, 分别存储在 content_features 和 style_features 字典中。

使用字典推导式遍历 style_features 字典中的每一层,对每一层的特征张量计算 Gram 矩阵,并将结果存储在 style_grams 字典中。

在字典推导式中,对于每个 layer,调用 gram_matrix 函数计算其对应的 Gram 矩阵,将结果存储在 style_grams 字典中,键为层名,值为对应层的 Gram 矩阵。

1
2
3
4
5
6
7
8
content_features = get_features(img=content, model=vgg)
style_features = get_features(img=style, model=vgg)
''' 使用字典推导式,遍历 style_features 字典中的每一层,
对每一层的特征张量计算 Gram 矩阵,并将结果存储在 style_grams 字典中,
键为层名,值为对应层的 Gram 矩阵'''
style_grams = {
layer: gram_matrix(tensor=style_features[layer]) for layer in style_features
}

代码实现:定义优化目标与权重

target = content.clone().requires_grad_(True): 创建了一个与内容图像 content 相同的张量 target,并设置其可导,用于后续的优化过程。

style_weights: 定义了不同层的权重,指定了每个层对风格损失的贡献权重。

content_weight = 1: 内容损失的权重,用于控制内容重建的贡献程度。

style_weight = 1e4: 风格损失的权重,用于控制风格重建的贡献程度。通常风格损失的权重会设置得比较大,以确保风格特征能够有效地传递到生成的图像中。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
# 定义优化目标
target = content.clone().requires_grad_(True)

# 定义不同层的权重
style_weights = {
"conv1_1": 1,
"conv2_1": 0.8,
"conv3_1": 0.5,
"conv4_1": 0.3,
"conv5_1": 0.1,
}

# 定义两种损失对应的权重
content_weight = 1
style_weight =1e4

训练模型

在训练模型进行风格迁移时,我们不断抽取合成图像的内容特征和风格特征,然后计算损失函数。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
# 开始训练
# 每 20% 输出一次训练结果
show_every = 2000
epoch = 20000
optimizer = optim.Adam([target], lr=0.003)

for i in tqdm(range(epoch)):
target_features = get_features(img=target, model=vgg)
# 计算内容损失,即目标图像与内容图像在特定层的特征之间的均方误差
content_loss = torch.mean(
(target_features["conv4_2"] - content_features["conv4_2"]) ** 2
)

style_loss = 0
for layer in style_weights:
target_feature = target_features[layer]
target_gram = gram_matrix(tensor=target_feature)
_, d, h, w = target_feature.shape
style_gram = style_grams[layer]
# 计算风格损失,遍历每一层的权重 style_weights,计算目标图像在该层的 Gram 矩阵与风格图像在该层的 Gram 矩阵之间的均方误差,并加权求和
layer_style_loss = style_weights[layer] * torch.mean(
(target_gram - style_gram) ** 2
)
style_loss += layer_style_loss / (d * h * w)

# 计算总损失,包括内容损失和风格损失,其中通过 content_weight 和 style_weight 分别控制两者的权重
total_loss = content_weight * content_loss + style_weight * style_loss

optimizer.zero_grad() # 使用优化器进行梯度清零
total_loss.backward() # 反向传播总损失
optimizer.step() # 更新参数

# 展示训练过程
if i % show_every == 0:
print("Total Loss:", total_loss.item())
plt.imshow(im_convert(tensor=target))
plt.show()

使用内容图像和风格图像来生成一个新的目标图像,使其同时保留内容和风格特征。具体步骤如下:

  1. 设置训练过程中的输出频率和总迭代次数。

  2. 定义优化器为 Adam,并传入目标图像 target 作为要优化的参数。

  3. 在每次迭代中,计算目标图像在不同层的特征表示,然后分别计算内容损失和风格损失。

  4. 内容损失计算目标图像与内容图像在特定层特征之间的均方误差。

  5. 风格损失通过计算目标图像在每层的Gram矩阵与风格图像在对应层的 Gram 矩阵的均方误差,并加权求和得到。

  6. 总损失由内容损失和风格损失加权求和得到。

  7. 使用优化器进行梯度清零,反向传播总损失,更新参数。

  8. 每隔一定迭代次数展示训练过程中的目标图像及当前总损失。

  9. 这个训练过程将逐步调整目标图像,使其在内容和风格上与内容图像和风格图像相似。通过迭代训练,最终可以得到一个同时保留了内容和风格特征的合成图像。

如下图所示,我们可以看到,合成图像保留了内容图像的风景和物体,并同时迁移了风格图像的色彩。 例如,合成图像具有与风格图像中一样的色彩块,其中一些甚至具有画笔笔触的细微纹理。

风格图像生成过程

用预训练的卷积神经网络完成色彩风格迁移
http://example.com/2024/05/04/styleTransfer/
作者
Guoming Huang
发布于
2024年5月4日
更新于
2024年5月5日
许可协议