运用pytorch实现基于VGG19的风格迁移（绝对能跑）

置顶 WHU菜菜子

已于 2023-12-18 11:08:01 修改

阅读量2.5k

点赞数 9

文章标签： pytorch 人工智能 python

于 2023-09-10 16:30:35 首次发布

本文链接：https://blog.youkuaiyun.com/U20215153/article/details/132792597

版权

注意：请先转化为JPG格式！确保安装有torch框架
完整版代码在最下边

style：

content：

output：

图像的读取与基本操作：

正文：

1.图像读取与大小调整

def load_image(image_path, transform=None, max_size=None, shape=None):
    image = Image.open(image_path)
    if max_size:
        scale = max_size / max(image.size)
        size = np.array(image.size) * scale
        image = image.resize(size.astype(int), Image.ANTIALIAS)

    if shape:
        image = image.resize(shape, Image.LANCZOS)

    if transform:
        image = transform(image).unsqueeze(0)

    return image.to(device)


transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406],
                         std=[0.229, 0.224, 0.225])
])

使用PIL库的Image.open函数打开给定路径的图像文件，将其加载为PIL图像对象，并存储在变量image中。如果指定了max_size参数，则计算缩放比例以确保图像的最大边不超过max_size。防止图像过大，减少计算负担。根据计算出的缩放比例scale，调整图像的尺寸。这里使用Image.resize方法，使用Image.ANTIALIAS滤波器以提高图像质量。调整后的图像大小存储在变量size中。如果指定了shape参数，再次调整图像的尺寸，将其调整为指定的形状。这里使用Image.LANCZOS滤波器来进行调整。如果指定了transform参数，将图像应用该变换。通常，这个变换将图像转换为PyTorch张量，并可能进行其他的预处理。转换后的图像被包装在一个大小为1的批次中（unsqueeze(0)），因为风格迁移需要处理批次的数据。最后，返回经过处理的图像，并将其移到计算设备（GPU或CPU）上，根据程序开头的设备选择。
2.变为张量类型并归一化

transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406],
                         std=[0.229, 0.224, 0.225])
])

transforms.Compose：这是PyTorch中的一个变换组合对象。它允许将多个图像变换按顺序组合在一起，以便在单个变换中一次性应用它们。transforms.ToTensor()：一个图像变换，将PIL图像或NumPy数组转换为PyTorch张量。它将图像的像素值缩放到0到1之间，并重新排列通道顺序。transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])：从每个通道的像素值中减去均值（mean）。然后，将结果除以标准差（std）。这个特定的均值和标准差来自于ImageNet数据集的统计信息，它们被广泛用于深度学习中的图像预处理。通过执行这个归一化操作，图像的像素值将被重新缩放，使其具有零均值和单位方差。
二.VGGNet类：

class VGGNet(nn.Module):
    def __init__(self):
        super(VGGNet, self).__init__()
        self.select = ['0', '5', '10', '19', '28']
        self.vgg = models.vgg19(pretrained=True).features

    def forward(self, x):
        features = []
        for name, layer in self.vgg._modules.items():
            x = layer(x)
            if name in self.select:
                features.append(x)
        return features

.
VGGNet 类继承自 nn.Module，这是 PyTorch 中构建神经网络模型的基类。super(VGGNet, self).__init__()：调用父类 nn.Module 的初始化方法，确保正确初始化子类。self.select：这是一个包含字符串的列表，指定了在 VGG 模型中要提取特征的层的名称。这些层的索引是基于 VGG19 模型的层的顺序。self.vgg：这是一个预训练的 VGG19 模型，使用 models.vgg19(pretrained=True).features 来加载。.features 表示只加载 VGG19 模型的特征提取部分，而不包括分类器部分。forward 方法：这是一个必须实现的方法，用于定义模型的前向传播过程。在这里，执行以下操作：features：这是一个空列表，用于存储提取的特征。for name, layer in self.vgg._modules.items()：迭代 VGG 模型的各个层。x = layer(x)：将输入 x 传递给当前层 layer，以计算特征。在每个前向传播步骤中，x 将经过一层的计算。if name in self.select：如果当前层的名称在 self.select 列表中，表示这是要提取特征的层之一，将 x 添加到 features 列表中。最后，返回 features 列表，其中包含了在 self.select 中指定的层的特征。

三．学习与优化过程：

1.初始化目标图像和优化器：

target = content.clone().requires_grad_(True)
optimizer = torch.optim.Adam([target], lr=0.11, betas=[0.5, 0.999])
vgg = VGGNet().to(device).eval()
target_features = vgg(target)
total_step = 10
style_weight = 400.

target：目标图像是从内容图像克隆而来的，它是一个 PyTorch 张量，具有与内容图像相同的初始值。通过 .clone() 方法，创建了一个内容图像的副本。requires_grad_(True)：这一步设置了 target 张量需要计算梯度。这是因为在风格迁移中，我们将调整 target 图像的像素值，以最小化损失函数，因此需要计算相对于 target 的梯度，以便在优化中更新图像。optimizer：这里使用了 Adam 优化器，用于优化目标图像以最小化损失。
2.风格迁移主循环：

for step in range(total_step):
    target_features = vgg(target)
    content_features = vgg(content)
    style_features = vgg(style)
    style_loss = 0
    content_loss = 0
    for f1, f2, f3 in zip(target_features, content_features, style_features):
        content_loss += torch.mean((f1 - f2) ** 2)
        _, c, h, w = f1.size()
        f1 = f1.view(c, h * w)
        f3 = f3.view(c, h * w)

        # 计算gram matrix
        f1 = torch.mm(f1, f1.t())
        f3 = torch.mm(f3, f3.t())
        style_loss += torch.mean((f1 - f3) ** 2) / (c * h * w)

    loss = content_loss + style_weight * style_loss

    # 更新target
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

    if step % 10 == 0:
        print("Step [{}/{}], Content Loss: {:.4f}, Style Loss: {:.4f}"
              .format(step, total_step, content_loss.item(), style_loss.item()))

style_weight：这是用于加权风格损失的超参数。风格损失的权重设置为 style_weight * style_loss，以便在总损失中平衡内容损失和风格损失的贡献。
for step in range(total_step)：这是一个训练循环，total_step 是训练的总迭代次数。
target_features = vgg(target)、content_features = vgg(content)、style_features = vgg(style)：这些行分别提取了目标图像、内容图像和风格图像的特征。这是通过传递图像到预训练的 VGG 模型 (vgg) 中来实现的。
style_loss 和 content_loss：这两个变量分别用于累积风格损失和内容损失。这两种损失在每个训练步骤中计算。for f1, f2, f3 in zip(target_features, content_features, style_features)：这是一个循环，用于遍历目标、内容和风格图像的特征。content_loss 的计算：通过计算特征之间的均方误差来衡量目标图像与内容图像的相似性。style_loss 的计算：通过计算特征之间的 Gram 矩阵的均方误差来衡量目标图像与风格图像的相似性。这是通过将特征张量的形状调整并计算 Gram 矩阵来实现的。
loss：总损失是内容损失和加权的风格损失之和。
optimizer.zero_grad()：清零优化器的梯度，以准备计算新的梯度。
loss.backward()：计算损失相对于目标图像的梯度。
optimizer.step()：根据梯度更新目标图像，以最小化总损失。

denormalize = transforms.Compose([transforms.Normalize(mean=[-0.485 / 0.229, -0.456 / 0.224, -0.406 / 0.225],
                                                       std=[1 / 0.229, 1 / 0.224, 1 / 0.225])])
img = denormalize(target.clone()).squeeze().cpu().detach().clamp_(0, 1)
img_pil = transforms.ToPILImage()(img)
img_pil.save("output_image.jpg")

四.图像反归一化和保存：

denorm 变换实际上是将图像反向归一化到原始像素范围的操作。这是因为在风格迁移中，目标图像在训练过程中经过了归一化，但在显示时需要将其还原到正常范围
这个变换的作用是将图像还原为原始像素范围。target.clone()：首先克隆了目标图像，以确保不修改原始数据。squeeze()：如果图像有一个额外的批次维度，这个操作会将其去除，因为在处理单个图像时，批次维度不再需要。cpu().detach()：将图像数据从 GPU 移到 CPU，并分离图像数据，使其不再与计算图相关联。clamp_(0, 1)：将图像数据截断在0到1之间，以确保像素值在合法范围内。这行代码使用 transforms.ToPILImage() 将反归一化后的图像数据 img 转换为 PIL 图像对象 img_pil。这是因为通常在显示、保存或进一步处理图像时，需要将图像数据转换为 PIL 图像对象

总代码：

from __future__ import division

import numpy as np
import torch
import torch.nn as nn
from PIL import Image
from torchvision import models
from torchvision import transforms

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(device)


def load_image(image_path, transform=None, max_size=None, shape=None):
    image = Image.open(image_path)
    if max_size:
        scale = max_size / max(image.size)
        size = np.array(image.size) * scale
        image = image.resize(size.astype(int), Image.ANTIALIAS)

    if shape:
        image = image.resize(shape, Image.LANCZOS)

    if transform:
        image = transform(image).unsqueeze(0)

    return image.to(device)


transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406],
                         std=[0.229, 0.224, 0.225])
])

content = load_image("C:/Users/qwe/Desktop/content666.jpg", transform, max_size=700)
style = load_image("C:/Users/qwe/Desktop/style.jpg", transform, shape=[content.size(2), content.size(3)])


class VGGNet(nn.Module):
    def __init__(self):
        super(VGGNet, self).__init__()
        self.select = ['0', '5', '10', '19', '28']
        self.vgg = models.vgg19(pretrained=True).features

    def forward(self, x):
        features = []
        for name, layer in self.vgg._modules.items():
            x = layer(x)
            if name in self.select:
                features.append(x)
        return features


target = content.clone().requires_grad_(True)
optimizer = torch.optim.Adam([target], lr=0.11, betas=[0.5, 0.999])
vgg = VGGNet().to(device).eval()
target_features = vgg(target)
total_step = 10
style_weight = 400.
for step in range(total_step):
    target_features = vgg(target)
    content_features = vgg(content)
    style_features = vgg(style)
    style_loss = 0
    content_loss = 0
    for f1, f2, f3 in zip(target_features, content_features, style_features):
        content_loss += torch.mean((f1 - f2) ** 2)
        _, c, h, w = f1.size()
        f1 = f1.view(c, h * w)
        f3 = f3.view(c, h * w)

        # 计算gram matrix
        f1 = torch.mm(f1, f1.t())
        f3 = torch.mm(f3, f3.t())
        style_loss += torch.mean((f1 - f3) ** 2) / (c * h * w)

    loss = content_loss + style_weight * style_loss

    # 更新target
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

    if step % 10 == 0:
        print("Step [{}/{}], Content Loss: {:.4f}, Style Loss: {:.4f}"
              .format(step, total_step, content_loss.item(), style_loss.item()))

denormalize = transforms.Compose([transforms.Normalize(mean=[-0.485 / 0.229, -0.456 / 0.224, -0.406 / 0.225],
                                                       std=[1 / 0.229, 1 / 0.224, 1 / 0.225])])
img = denormalize(target.clone()).squeeze().cpu().detach().clamp_(0, 1)
img_pil = transforms.ToPILImage()(img)
img_pil.save("output_image.jpg")