训练PyTorch CNN实现图像风格迁移并适配C++部署的方案咨询

为什么你修改优化器后程序无法运行 #

你参考的官方教程属于图像像素优化式风格迁移，核心逻辑是固定预训练VGG模型（冻结所有参数，无训练空间），通过迭代优化生成图像的像素值来匹配风格，而非训练一个可复用的生成模型。原代码优化的是generated_image这个可训练张量，改成优化model.parameters()时，因模型没有可训练参数，自然会报错。

正确的方向：训练生成式风格转换模型 #

要得到可部署到C++的模型，需要训练端到端的生成器CNN——输入任意图像直接输出风格化结果，训练完成后可导出为TorchScript格式部署。主流方案包括：

• 快速风格迁移（Fast Neural Style Transfer）：专为实时风格转换设计，训练时用预训练VGG提取特征计算内容损失+风格损失，优化生成器的参数
• CycleGAN/Pix2Pix：适合跨域风格转换（如照片转艺术风格），支持无配对数据训练

针对灰度风格转换的具体实现方案 #

如果你的目标是将彩色图像转为特定灰度风格（如手绘质感、复古灰度），可以用轻量生成器模型实现，以下是完整流程：

1. 定义生成器模型 #

采用下采样+上采样的轻量结构，输入RGB图像，输出单通道灰度图：

import torch
import torch.nn as nn

class GrayStyleGenerator(nn.Module):
    def __init__(self):
        super().__init__()
        # 下采样卷积层
        self.down_layers = nn.Sequential(
            nn.Conv2d(3, 64, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(2),
            nn.Conv2d(64, 128, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(2)
        )
        # 上采样转置卷积层
        self.up_layers = nn.Sequential(
            nn.ConvTranspose2d(128, 64, kernel_size=2, stride=2),
            nn.ReLU(inplace=True),
            nn.ConvTranspose2d(64, 1, kernel_size=2, stride=2),
            nn.Sigmoid()  # 输出0-1范围的灰度图像
        )
    
    def forward(self, x):
        x = self.down_layers(x)
        x = self.up_layers(x)
        return x

2. 训练流程 #

训练时结合内容损失（保证生成图与原图结构一致）和风格损失（模仿目标灰度风格）：

import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
import torchvision.models as models

# 数据预处理
transform = transforms.Compose([
    transforms.Resize((256, 256)),
    transforms.ToTensor()
])

# 加载内容数据集（彩色图像）和风格数据集（目标灰度风格图像）
content_loader = DataLoader(
    datasets.ImageFolder(root="path/to/color_images", transform=transform),
    batch_size=4, shuffle=True
)
style_loader = DataLoader(
    datasets.ImageFolder(root="path/to/gray_style_images", transform=transform),
    batch_size=4, shuffle=True
)

# 设备配置
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
generator = GrayStyleGenerator().to(device)
optimizer = optim.Adam(generator.parameters(), lr=1e-4)
mse_loss = nn.MSELoss()

# 加载预训练VGG用于特征提取（冻结参数）
vgg = models.vgg19(pretrained=True).features.to(device).eval()
for param in vgg.parameters():
    param.requires_grad = False

# 特征提取函数
def extract_features(x, model):
    feature_layers = {"0": "conv1_1", "5": "conv2_1", "10": "conv3_1"}
    features = {}
    for name, layer in model._modules.items():
        x = layer(x)
        if name in feature_layers:
            features[feature_layers[name]] = x
            if name == "10":
                break
    return features

# Gram矩阵计算（用于风格损失）
def gram_matrix(x):
    _, channels, h, w = x.size()
    x_flat = x.view(channels, h * w)
    return torch.mm(x_flat, x_flat.t()) / (channels * h * w)

# 训练循环
num_epochs = 50
for epoch in range(num_epochs):
    generator.train()
    total_epoch_loss = 0.0
    for (content_imgs, _), (style_imgs, _) in zip(content_loader, style_loader):
        content_imgs = content_imgs.to(device)
        style_imgs = style_imgs.to(device)
        
        # 生成风格化图像
        generated_imgs = generator(content_imgs)
        # 将单通道生成图转为三通道，匹配VGG输入要求
        generated_rgb = generated_imgs.repeat(1, 3, 1, 1)
        
        # 提取特征
        content_feats = extract_features(content_imgs, vgg)
        generated_feats = extract_features(generated_rgb, vgg)
        style_feats = extract_features(style_imgs, vgg)
        
        # 计算内容损失
        content_loss = mse_loss(generated_feats["conv3_1"], content_feats["conv3_1"])
        
        # 计算风格损失
        style_loss = 0.0
        for layer in ["conv1_1", "conv2_1", "conv3_1"]:
            gen_gram = gram_matrix(generated_feats[layer])
            style_gram = gram_matrix(style_feats[layer])
            style_loss += mse_loss(gen_gram, style_gram)
        
        # 总损失（调整风格损失权重）
        total_loss = content_loss + 1e5 * style_loss
        
        # 反向传播
        optimizer.zero_grad()
        total_loss.backward()
        optimizer.step()
        
        total_epoch_loss += total_loss.item()
    
    print(f"Epoch {epoch+1}/{num_epochs}, Average Loss: {total_epoch_loss/len(content_loader):.4f}")

# 保存模型权重
torch.save(generator.state_dict(), "gray_style_generator.pth")

3. 模型部署 #

将训练好的模型转为TorchScript格式，即可部署到C++：

generator.eval()
# 用示例输入追踪模型
example_input = torch.randn(1, 3, 256, 256).to(device)
traced_model = torch.jit.trace(generator, example_input)
# 保存TorchScript模型
traced_model.save("gray_style_generator.pt")

可用的PyTorch预训练风格转换模型 #

• PyTorch Hub提供预训练的快速风格迁移模型，可直接加载使用
• 开源项目fast-neural-style包含多种风格的预训练模型，可直接下载微调
• CycleGAN的PyTorch实现仓库提供多种跨域风格迁移的预训练模型，包含灰度风格相关的预训练权重

内容的提问来源于stack exchange，提问作者Rreit