图像风格迁移代码(PyTorch)

以下是一个简单的图像风格迁移代码，使用PyTorch实现：

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import torchvision.transforms as transforms
import torchvision.models as models
import torchvision.datasets as datasets
import matplotlib.pyplot as plt

# 定义图像处理的转换
transform = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

# 加载图像数据集
content_dataset = datasets.ImageFolder('content_images_folder_path', transform=transform)
style_dataset = datasets.ImageFolder('style_images_folder_path', transform=transform)

# 创建数据加载器
content_loader = torch.utils.data.DataLoader(content_dataset, batch_size=1, shuffle=True)
style_loader = torch.utils.data.DataLoader(style_dataset, batch_size=1, shuffle=True)

# 加载预训练的VGG19模型
vgg = models.vgg19(pretrained=True).features
for param in vgg.parameters():
    param.requires_grad_(False)

# 将模型移至GPU
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
vgg.to(device)

# 定义内容损失函数
def content_loss(content_features, target_features):
    return torch.mean((content_features - target_features)**2)

# 定义风格损失函数
def gram_matrix(input):
    a, b, c, d = input.size()
    features = input.view(a * b, c * d)
    G = torch.mm(features, features.t())
    return G.div(a * b * c * d)
    
def style_loss(style_features, target_features):
    style_gram = gram_matrix(style_features)
    target_gram = gram_matrix(target_features)
    return torch.mean((style_gram - target_gram)**2)

# 定义优化器
optimizer = optim.Adam([target], lr=0.01)

# 迭代次数
epochs = 500

# 定义目标图像
target = torch.randn(content_dataset[0][0].shape, requires_grad=True).to(device)

# 迭代训练
for epoch in range(epochs):
    for content, style in zip(content_loader, style_loader):
        # 将内容图像和风格图像传入模型
        content_features = vgg(content[0].to(device))
        style_features = vgg(style[0].to(device))
        target_features = vgg(target)

        # 计算损失函数
        content_loss_value = content_loss(content_features, target_features)
        style_loss_value = style_loss(style_features, target_features)
        loss = content_loss_value + style_loss_value

        # 反向传播优化
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        # 打印损失函数和目标图像
        if epoch % 100 == 0:
            print(f"Epoch {epoch}, Content Loss: {content_loss_value.item()}, Style Loss: {style_loss_value.item()}")
            plt.imshow(target.detach().cpu().squeeze().permute(1, 2, 0).numpy())
            plt.show()

需要将代码中的`content_images_folder_path`和`style_images_folder_path`修改为实际的内容图像和风格图像所在的文件夹路径。