ARTICLE DETAIL

资讯详情

深耕郑州网站建设与运营推广的一线实战洞察。

基于VGG16的图像风格迁移实战:从PyTorch训练到Flask部署

基于VGG16的图像风格迁移实战:从PyTorch训练到Flask部署 简介面向计算机相关专业正在做毕业设计、课程设计或期末大作业的学生以及其他需要项目实战练习的Python学习者这是一份基于CNN卷积神经网络的图像风格迁移完整项目源码。资源共93个文件包含Python训练/测试脚本、视频风格化工具、Web界面相关代码、预训练模型权重.pth、大量图片素材jpg/png以及演示视频mp4压缩包整体57.03MB目录结构清晰便于按模块阅读与复用。目前已有222人学习下载。项目提供了从模型训练、图片与视频风格化到交互式展示的完整链路代码完整可运行并带有可直接调用的训练权重既能快速体验将照片转换为星空、马赛克等艺术风格也适合在此基础上修改数据集与超参数开展二次实验有助于深入理解CNN特征提取、风格损失与图像重建机制是较有参考价值的高分毕业设计范例。1. 图像风格迁移背后的 CNN 特征解耦以及这套源码到底给了你什么第一次跑通风格迁移时很多人会愣一下为什么把梵高的星空喂给 VGG16输出的不是分类结果而是一张重绘的风景照因为卷积神经网络在逐层抽象的过程中浅层卷积核记住的是边缘、纹理、笔触深层记住的是物体结构。风格迁移的本质是把浅层特征的统计分布从风格图里提取出来再约束另一张内容图去匹配这个分布。这个项目做的就是这件事而且它把流程拆成了三部分离线训练风格模型、单张图片推理、基于 Flask 的视频和 Web 展示。对正在做毕业设计的人这套源码的价值不在于能跑而在于它把训练、推理、展示串成了一条完整的链路。checkpoints 目录里已经有starry_night_28000_vgg16.pth、mosaic_10000.pth、cuphead_10000.pth、sketch_2000.pth四个预训练权重也就是你不需要 GPU 也能立刻拿本就训练好的模型做迁移。而对想自己训练风格的人来说train.py和make_style_new_dataset.py提供了从原始风格图到可复现训练集的完整路径。接下来我会按推理链路、训练配置、Web 部署、特征层调优的顺序把这个项目的每个螺丝都拆开看一遍。2. 先看懂推理链路test_on_image.py 与 VGG16 前向传播2.1 特征提取与 Gram 矩阵风格迁移的数学底座在动手调用代码之前得先理解模型文件里存的是什么。VGG16 在 ImageNet 上做分类训练全连接层之前的卷积部分本质上是把图像编码成一摞特征图。风格迁移的做法是让内容图经过 VGG16取出中间某一层的特征图让风格图也经过 VGG16但取的是多个层级的特征图并计算它们的 Gram 矩阵。这个矩阵统计的是通道之间的相关性说白了就是在描述这种纹理通常会搭配哪种纹理这正是风格的数学定义。def gram_matrix(feature_map): b, c, h, w feature_map.shape features feature_map.view(b * c, h * w) gram torch.mm(features, features.t()) return gram / (c * h * w)这段代码计算的是某一张特征图内部的 Gram 矩阵把通道维和空间维压平再用矩阵乘法算出通道间的内积最后除以通道数和像素数做归一化。在项目推理脚本里这个函数会被反复调用——每提取一层风格特征就要算一次 Gram 矩阵。归一化的意义在于不同尺度的图产生的内积数值量级差异很大不归一化的话浅层特征对损失的贡献会被数值淹掉。2.2 加载预训练权重开始单图迁移项目根目录下的test_on_image.py是推理入口。它的核心逻辑并不复杂先加载一张内容图加载一个训练好的.pth权重文件然后用这个权重直接对内容图的像素做迭代优化。注意推理过程中模型的参数是被冻结的训练过程改变的不是 VGG16 的权重而是输入图像的像素值。python test_on_image.py \ --content_image images/1111.jpg \ --style_image styles/mosaic.jpg \ --checkpoint checkpoints/mosaic_10000.pth \ --output_image stylized-777.jpg \ --cuda 0参数含义是这样的--content_image和--style_image分别指定内容和风格图来源虽然风格图在推理阶段的意义已经不大风格信息已经固化在.pth里了但脚本仍然要求传入以确保数据流一致--checkpoint指向预训练权重--output_image定义迁移结果保存路径--cuda 0指定使用第一块 GPU电脑没有 NVIDIA 显卡就把这个参数去掉脚本会退回到 CPU 推理模式。这里有一点值得注意.pth文件里保存的并不是 VGG16 的完整状态字典而是训练出来的一个生成网络或者叫 Transformer Network。也就是说在models.py里定义了一个轻量的卷积生成器它接收普通照片输入直接输出风格化后的图像。推理时做一次前向就完成了不需要像传统风格迁移那样反复迭代几百步。2.3 视频推理test_on_video.py 的逐帧处理与显存管理test_on_video.py把单图推理扩展到了视频场景。它的做法很朴素用 OpenCV 打开视频流逐帧读取每帧经过生成网络得到风格化结果再写进输出视频。但朴素不等于简单里面有三个工程细节值得学习。cap cv2.VideoCapture(input_video) fps int(cap.get(cv2.CAP_PROP_FPS)) width int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) height int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) out cv2.VideoWriter(output_video, cv2.VideoWriter_fourcc(*mp4v), fps, (width, height)) while True: ret, frame cap.read() if not ret: break tensor_frame preprocess_frame(frame) with torch.no_grad(): stylized generator(tensor_frame) output_frame postprocess_frame(stylized) out.write(output_frame)视频推理的关键问题在于显存一张 1080p 的帧经过生成网络中间特征图的体积会非常大如果直接前向会爆显存。常见的处理方式是把帧缩放到 512 或 640 边长再做推理然后再把输出拉伸回原始分辨率。这个项目里preprocess_frame函数内部做的就是 resize、BGR 转 RGB、归一化、转 Tensor而postprocess_frame做的是逆操作。如果你的视频卡顿优先检查是不是前处理时没有做缩放。参数说明建议值--content_size内容图输入边长512--style_size风格图缩放边长训练用256--checkpoint模型权重路径checkpoints 目录下 .pth--cuda指定 GPU 设备号无 GPU 则省略--output_video结果视频保存路径.mp4 格式这段表格里的--content_size和--style_size是这类项目里最高频的两个超参数后面训练章节还会再提到它们。你已经理解了推理链路下一步要搞清楚的是这些预训练权重是怎么来的训练流程又该如何配置3. 训练自己的风格模型从风格数据集到可复用的 .pth 权重3.1 制作训练数据集make_style_new_dataset.py 在干什么项目里的make_style_new_dataset.py解决的问题非常实际训练一个风格迁移模型不能只靠一两张图。如果只把一张风格图拿来反复训练模型会过拟合到这张图的色彩分布上换一张内容图结果可能惨不忍睹。所以要做的是把一张风格图拆解出多个尺度的子图或者从内容数据集里采样一批训练图让模型看到足够多的内容多样性。python make_style_new_dataset.py \ --style_image styles/monet.jpeg \ --content_dir images/ \ --output_dir content_dataset/ \ --num_samples 5000 \ --crop_size 256这段脚本的逻辑是从--content_dir里的内容图片中随机裁剪256x256区域每张图裁剪若干次凑足 5000 个训练样本输出到content_dataset/目录。风格图本身不参与裁剪它只是作为风格损失的目标。这么做的好处是模型在学习风格时内容的多样性足够不会把某张特定照片的布局背下来。3.2 训练主循环与损失函数配置train.py是整套源码的核心。它定义了一个生成网络 G一个固定的 VGG16 特征提取器以及内容和风格两个损失项。训练时生成网络接收内容图输出风格化图风格化图和内容图分别过 VGG16 提取特征计算内容损失风格化图和风格图分别提取多尺度特征计算风格损失的加权和。style_loss 0 for layer, weight in style_layers.items(): gen_feat vgg16(stylized)[layer] style_feat vgg16(style_img)[layer] style_loss weight * mse_loss(gram_matrix(gen_feat), gram_matrix(style_feat)) content_loss mse_loss(vgg16(stylized)[relu3_3], vgg16(content_img)[relu3_3]) total_loss 1e4 * style_loss 1e0 * content_loss total_loss.backward() optimizer.step()内容损失选用的是 VGG16 的relu3_3层特征风格损失则跨了relu1_1、relu2_1、relu3_1、relu4_1、relu5_1五个层。这个选择不是随意的内容损失如果选太浅的层重建出来的是像素级的吻合风格化效果会很弱选太深的层物体的语义结构保得住但边缘会被抹掉。relu3_3是语义和纹理的中间地带。风格损失用多层加权是为了让不同尺度的纹理模式都被捕捉到。值得注意损失权重风格损失权重是内容损失的 10000 倍这不是模型作者拍脑袋定的而是因为风格损失的数值量级远小于内容损失需要大幅放大才能让风格信息在梯度更新中占据主导。如果你想让迁移结果更接近风格图的色彩把1e4调大到5e4如果想让内容结构更清晰把1e0调大到5e0。比较快的调参路径是先固定风格权重只动内容权重。3.3 训练收敛判断与 checkpoint 命名规律项目 checkpoint 目录里starry_night_28000_vgg16.pth这个名字暴露了训练配置28000是迭代步数vgg16是特征提取器类型。训练完以后模型保存的不只是权重还应该在models.py里定义了完整的生成网络结构加载时用torch.load读到的是生成网络整个实例的 state_dict。checkpoint torch.load(checkpoints/starry_night_28000_vgg16.pth, map_locationcpu) generator TransformerNetwork() generator.load_state_dict(checkpoint[model_state_dict]) generator.eval()训练什么时候可以停判断标准很简单把验证集里几张内容图固定下来每 500 步跑一次推理肉眼观察风格化结果。风格化结果的纹理结构稳定了、不再跟着迭代明显变化就可以停了。28000 步这个量级在单张 2080Ti 上大约需要 2 小时所以训练时记得用--checkpoint_interval 2000自动保存中间权重这样即使中断也能从最近的权重续训。到这里你已经能自己训练一个风格模型了。但拿train.py训练出来的生成网络去做推理和项目里已经预训练的starry_night_28000_vgg16.pth之间还有一层 Web 封装。下一节就讲这个怎么把模型服务化让别人能在浏览器里玩。4. 把模型封装成 Web 服务app.py 的 Flask 路由与前端交互4.1 Flask 接收图片并返回风格化结果项目根目录的app.py用 Flask 搭建了一个 Web 演示入口。这个文件的价值在于它展示了如何把 PyTorch 模型嵌入到一个 Web 请求的生命周期里请求进来、读图、前向、返回结果、释放资源。这里最关键的设计是全局初始化模型而不是每次请求都重新torch.load——模型加载的 IO 开销远大于一次前向如果放在请求处理函数里并发一上来进程就直接卡死。app Flask(__name__) def load_model(): global generator checkpoint torch.load(checkpoints/mosaic_10000.pth, map_locationcpu) generator TransformerNetwork() generator.load_state_dict(checkpoint[model_state_dict]) generator.eval() app.route(/stylize, methods[POST]) def stylize(): file request.files[image] img_bytes np.frombuffer(file.read(), np.uint8) img cv2.imdecode(img_bytes, cv2.IMREAD_COLOR) tensor preprocess_for_generator(img) with torch.no_grad(): out generator(tensor) result_img postprocess_to_bgr(out) _, encoded cv2.imencode(.jpg, result_img) return Response(encoded.tobytes(), mimetypeimage/jpeg)注意这里用了map_locationcpu即使服务器没有 GPU也能正常跑推理。但如果你有 GPU建议改成map_locationcuda:0否则每个请求都会在 CPU 上做前向RTX 3060 跑一次 VGG16 风格迁移大约需要 2 秒CPU 可能要 20 秒以上。4.2 前端模板与静态资源的配合方式templates/和static/两个目录的职责非常清晰。templates/index.html是主页面包含一个文件上传控件、一个展示区域和一段简单的 JavaScript用fetch把文件 POST 到/stylize接口然后接收返回的 JPEG 数据渲染到img标签里。static/下放的是 css 和若干示例效果图比如stylized-777.jpg、stylized-555.jpg这些图直接作为展示案例写死在页面上方便用户还没上传时就对效果有个预期。input typefile idupload acceptimage/* button onclickuploadImage()风格化/button img idresult alt结果图这段前端的核心简化版就三行上传控件拿到本地文件按钮触发uploadImage()结果图直接替换img标签。真正的uploadImage()函数里需要用FormData构造请求体fetch(/stylize, { method: POST, body: formData })然后把response.blob()转成URL.createObjectURL塞给img标签。这个模式值得抄进你自己的课程设计展示页里。4.3 启动 Web 服务与排错备忘python app.py # Flask 默认监听 5000 端口 # 浏览器访问 http://localhost:5000启动之前先确认三件事templates/index.html必须存在Flask 找不到模板会直接抛TemplateNotFoundcheckpoints/下的.pth文件名要用os.path.exists检查一遍路径写错了会在load_model()阶段报FileNotFoundError如果你的 Python 版本是 3.9但 models 目录里的.pyc是cpython-37或cpython-310的这不影响运行——.pyc只会在 Python 解释器版本不匹配时被忽略源文件会重新编译。如果你要用这个 Web 服务做局域网演示把启动参数加上--host 0.0.0.0这样同网段的其他电脑也能访问。这里还有一个 Python 版本兼容的坑checkpoints/目录下同时出现了models.cpython-37.pyc和models.cpython-310.pyc说明项目作者在不同的 Python 版本下各跑过一次。你用 Python 3.9 新环境时要重新跑一次脚本让解释器生成匹配的.pyc否则首次启动会稍慢一些。5. 特征层选择与损失权重把迁移效果从能看调到好看5.1 手动指定 VGG16 特征层权重训练脚本里style_layers是一个字典源码默认值大致是relu1_1: 1.0, relu2_1: 0.8, relu3_1: 0.4, relu4_1: 0.2, relu5_1: 0.1。浅层权重高、深层权重低这个递减配比是有道理的浅层捕捉笔触和细小纹理对风格的还原最重要深层捕捉物体级的大尺度构成如果权重偏高风格化图会出现奇怪的色块。但具体到不同风格配比需要相应调整。油画/印象派风格如starry_night.pth、monet.jpeg提高浅层relu1_1到 1.5保留更多笔触感。马赛克/像素化风格如mosaic.pth降低浅层权重到 0.5否则马赛克块会被细节纹理干扰。素描风格如sketch.pth提高relu4_1和relu5_1权重素描的大线条更依赖中深层特征。style_layers { relu1_1: 1.5, relu2_1: 0.8, relu3_1: 0.4, relu4_1: 0.3, relu5_1: 0.2 }修改后直接跑train.py不需要改任何其他代码。如果在同一个.pth上反复调参不理想建议你在checkpoints/里新建一个目录把不同配比训练的权重分开存放方便做横向对比实验。5.2 大图推理的显存分块技巧用训练好的权重对一张 4000×3000 的照片做推理时直接送进生成网络大概率会 OOM。源码里的--content_size参数只能在预处理时缩放整张图想要保留细节又要跑得动就需要按重叠块推理再拼接。这个技巧在很多风格迁移项目里都是通用的代码大致如下def inference_large_image(generator, large_img, patch_size512, overlap32): h, w large_img.shape[:2] output np.zeros_like(large_img) for y in range(0, h, patch_size - overlap): for x in range(0, w, patch_size - overlap): patch large_img[y:ypatch_size, x:xpatch_size] patch_tensor preprocess(patch) with torch.no_grad(): stylized generator(patch_tensor) result postprocess(stylized) output[y:ypatch_size, x:xpatch_size] result return output每块 512×512重叠 32 像素是为了避免相邻块的拼接痕迹。如果你的风格模型很吃上下文比如粗笔触油画重叠区域建议扩大到 64 像素否则块与块之间会有一道肉眼可见的接缝。注意这段代码在边界处没有做 Alpha 混合真正工程化时可以在重叠区域做线性插值视觉上会自然很多。5.3 用预训练权重快速验证风格强度你拿到项目后不需要重新训练就能验证不同参数对结果的影响。用test_on_image.py加载cuphead_10000.pth然后对同一张内容图分别加--style_size 128、256、512跑三次对比输出图的纹理密度——style_size越小风格纹理越粗犷越大风格纹理越细腻。这是因为训练时风格图的缩放尺度决定了模型见到的纹理颗粒大小。这只影响推理时preprocess和postprocess的实现方式不影响weights本身。排查结果不理想时优先检查三个位置内容图是否过曝高光区域容易产生伪影、--cuda参数与当前环境是否匹配、.pth文件名是否与models.py里的类定义对应上。把generator.eval()加上再进行下一步调试这批权重在.eval()模式下的输出最稳定。本文还有配套的精品资源点击获取
返回列表