图像重建新视角：CNN+混合注意力如何破解去水印难题

简介：这是百度网盘AI大赛去水印模型冲刺赛的冠军方案完整代码与文档包，定位清晰，主要面向从事图像生成恢复、低层视觉任务研究的开发者、科研人员，以及准备参加同类AI比赛的选手。方案针对从带水印图像恢复原始图像这一低层次生成任务，基于CNN构建去水印模型，并通过混合注意力机制增强特征表达，同时采用多类数据增强策略防止过拟合、提升泛化能力。内容覆盖数据准备、网络结构设计、训练策略到推理部署的关键环节，工程化程度较高，适合作为论文复现或实际项目落地的参考实现。压缩包共166个文件，以103个Python脚本为核心，同时包含57个编译后的pyc文件、2个工程配置文件、1个模型权重文件（.pd），以及README说明与LICENSE许可，整体大小约92.74MB；源码结构清晰，Python脚本涵盖模型定义、训练与评估逻辑，pyc文件便于直接调用，工程配置帮助快速导入开发环境，训练好的权重可离线加载使用，节省重训时间。目前已有377人学习浏览，可帮助快速理解冠军思路，并迁移应用到图像修复、去噪等相似视觉任务中。

1. 去水印模型的冠军方案：先想清楚这是一道重建题，不是抠图题

手头这张隔着玻璃窗拍的证件照，右下角正好压着一块半透明水印。很多人第一反应是“把水印区域涂掉”，但涂掉只会留下一块糊斑。真正能上比赛的方案，是把去水印当成低层次图像重建任务：输入带水印的图，输出没有水印的原始画面。水印盖住的不是空白，是纹理、颜色、边缘，模型要“补”的东西远比“擦”的东西多。这套百度网盘AI大赛去水印模型冲刺赛冠军方案，核心就是CNN主干加混合注意力，再配一套完整的数据增强管线。对正在做图像修复、照片增强、乃至OCR前处理的算法工程师来说，这份代码包的价值在于——它把“重建质量怎么提上去”和“训练怎么稳住”两个老大难问题，同时给了可复现的答案。

2. 选型摊开讲：CNN主干加混合注意力，为什么比赛里够打

2.1 为什么这时候还选CNN，Transformer不香吗

比赛场景和工业落地的差别，很多人一开始没意识到：比赛数据量有限，推理资源有上限，而且水印这个任务本身没有长距离依赖问题——水印是局部覆盖，恢复也主要靠局部邻域信息。Transformer的全局注意力在这种任务上的优势不明显，反而把显存和训练轮次拉高不少。

CNN的归纳偏置在这里是优势的：卷积天然假设邻近像素相关，对纹理生成友好。普通卷积的问题是感受野不够深、通道选择不够聪明，于是方案里在CNN主干上挂了混合注意力机制。这个“混合”指的不是把CNN和Transformer拼一起，而是把通道注意力和空间注意力并行叠加在一起，让网络既能决定“看哪些特征图”，又能决定“看哪些像素区域”。

我用一张表把这几个选型的区别列出来。这张表不是官方文档，是我自己拆的时候做的笔记，方便你权衡：

2.2 sa_aidr.py 里到底拆成了哪几段

sa_aidr.py 这个文件名拆开看，sa 是 self-attention 的缩写，aidr 可以理解为 image de-watermark removal。主体结构我按惯例拆成三截：浅层特征提取、深层特征变换、重建输出。这里给一个参考结构伪代码，用 PyTorch 写法示意，原仓库用 Paddle 实现，但模块逻辑是同一套：

import torch import torch.nn as nn class ChannelAttention(nn.Module): def __init__(self, channels, reduction=8): super().__init__() self.fc = nn.Sequential( nn.AdaptiveAvgPool2d(1), nn.Conv2d(channels, channels // reduction, kernel_size=1), nn.ReLU(inplace=True), nn.Conv2d(channels // reduction, channels, kernel_size=1), nn.Sigmoid() ) def forward(self, x): # 全局平均池化压成通道描述符，再走两个1x1卷积，得到每个通道的权重 scale = self.fc(x) return x * scale class SpatialAttention(nn.Module): def __init__(self, kernel_size=7): super().__init__() padding = kernel_size // 2 self.conv = nn.Conv2d(2, 1, kernel_size=kernel_size, padding=padding) def forward(self, x): # 在通道维上取均值池化 max 池化，拼成2通道图，再卷积出空间权重 avg_out = torch.mean(x, dim=1, keepdim=True) max_out, _ = torch.max(x, dim=1, keepdim=True) attn = torch.cat([avg_out, max_out], dim=1) attn = self.conv(attn) attn = torch.sigmoid(attn) return x * attn class MixedAttnBlock(nn.Module): def __init__(self, channels): super().__init__() self.channel_attn = ChannelAttention(channels) self.spatial_attn = SpatialAttention() self.conv = nn.Conv2d(channels, channels, kernel_size=3, stride=1, padding=1) def forward(self, x): # 先通道注意力，再空间注意力，最后带一条恒等映射防梯度消失 identity = x x = self.channel_attn(x) x = self.spatial_attn(x) x = self.conv(x) return x + identity

通道注意力这一段要注意的是，池化方式直接影响权重质量。方案里用平均池化为主、最大池化为辅——平均池化能反映全局响应，最大池化能保住最显著特征。两者在空间注意力里拼接成双通道再卷积，是经典做法。恒等映射这一段不能省，去水印网络深度一旦超过二十层，没有残差连接训练基本会翻车。

2.3 参数配置参考：复现的时候从这套起步

我根据权重文件名 v7ms2_ep49.pd 和源码结构，整理了一套大概率走得通的配置。这不是源码里明文写的唯一配置，而是这个场景下合格从业者最可能用的方案：

这组配置里最容易踩的坑是我后面要单独讲的：空间注意力卷积核别设成 1x1，否则它只会学到一个逐像素的权重，相当于没感知上下文，水印和背景纹理一旦相似就分不开。通道注意力的 reduction 也别压到 4 以下，参数会暴增，在 512x512 输入下显存很容易见底。

3. 数据处理管线：水印是“造”出来的，增强是“防”出来的

3.1 dataset2.py 的数据增广思路

很多做图像任务的人对数据增强的理解还停留在旋转、翻转、调亮度三件套。去水印任务比这麻烦——模型没见过足够多样的水印形态，它就会把“淡色文字”“浅色logo”这类东西也当成水印去抹，结果把原图细节毁掉。

dataset2.py 里最关键的一段，是模拟训练数据生成。真实场景中收集同一张图的带水印和干净版成对数据很困难，比赛和工业落地里最常见的做法是：用干净图作为底图，然后程序化生成水印叠上去。生成参数必须包括水印字体、字号、旋转角度、透明度、位置，还有颜色。

我按这个方法重写一版增强管线，核心逻辑如下：

import random import numpy as np from PIL import Image, ImageDraw, ImageFont, ImageFilter def synthesize_watermark(background, font_paths, alpha_range=(0.3, 0.8)): """ 在干净图上模拟生成一个水印叠加层 background: PIL Image, RGB font_paths: 候选字体路径列表 返回: 带水印图, 水印灰度mask, 叠加透明度 """ bg = background.convert("RGBA") overlay = Image.new("RGBA", bg.size, (0, 0, 0, 0)) draw = ImageDraw.Draw(overlay) # 随机从字体库挑一个，字体大小要匹配底图分辨率 font = ImageFont.truetype(random.choice(font_paths), size=random.randint(48, 96)) text = "".join(random.choices("ABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789", k=random.randint(4, 8))) text_w, text_h = draw.textbbox((0, 0), text, font=font)[2:] pos_x = random.randint(0, max(0, bg.size[0] - text_w)) pos_y = random.randint(0, max(0, bg.size[1] - text_h)) # alpha 随机，水印颜色偏白或偏灰，比赛常见水印大多是白字半透明 alpha = random.uniform(*alpha_range) color = (255, 255, 255, int(255 * alpha)) draw.text((pos_x, pos_y), text, font=font, fill=color) # 加一点模糊模拟真实水印与背景的融合边界 if random.random() > 0.5: overlay = overlay.filter(ImageFilter.GaussianBlur(radius=1)) # mask 用来监督注意力或者给loss加权，二值化水印区域 mask = overlay.split()[3].point(lambda p: 255 if p > 0 else 0) merged = Image.alpha_composite(bg, overlay).convert("RGB") return merged, mask

这段里容易被忽略的参数是 alpha_range。我见过有人把透明度范围设成 0.1 到 0.9，结果模型在低透明度时根本看不到水印却强行输出“干净图”，把不该动的纹理也糊掉了。一般我会收敛到 0.3～0.8，兼顾可见度和难度。字体路径集合至少要准备三到五套不同风格字体，因为水印的笔画密度差异会让CNN学到完全不同的特征。mask 输出很多人嫌麻烦不存，但后续如果你想在水印区域上加强loss权重，或者想引导注意力模块聚焦，这个mask会非常有用。

3.2 防过拟合的增强组合：不是越多越好

数据增强有个反直觉的现象：在去水印任务上，过强的几何增强反而会毁掉重建质量。比如随机旋转超过30度，水印区域的位置和方向变化太剧烈，模型被迫去学习“旋转不变性”而不是“水印去除”。方案里采用的几种增强，我按作用拆成两部分。

第一类是刚体变换，概率控制在 0.5 左右：随机水平翻转、随机旋转 15 度以内、随机裁剪到 512x512 并缩放回原尺寸。第二类是光度扰动，这是防止模型偷懒的关键：随机调整亮度对比度饱和度、加入少量高斯噪声、偶尔叠加一次随机光照渐变。光度扰动让模型不能仅靠像素均值判断“这里是不是水印”，而被迫去学纹理层面的一致性。

代码里比较关键的是顺序：先做随机裁剪和翻转，再做水印合成，最后做光度扰动。顺序反了会导致水印和背景一起被调亮调暗，模型学到的关联性就会变弱。我复盘自己翻车经历时发现，数据管线顺序的影响比想象中大得多，同样一个增强集合，换个顺序验证集PSNR能差 0.4 以上。

4. 训练稳定态：loss 怎么配、学习率怎么退、权重文件里藏着什么

4.1 损失函数的组合逻辑：不能只盯像素差

去水印最基础的loss是L1和L2。L2对大误差惩罚重，但容易把结果磨得过于平滑；L1保边缘好，但收敛慢而且对噪声敏感。方案中的混合注意力给足了表达力，loss 配比没跟上的话，模型依然会得出水印没了、纹理也没了的“干净废图”。

我通常采用的做法是 L1 和 L2 各占一半，再叠加一个可选的感知loss。感知loss不是必需项，但对这种纹理重建任务帮助明显，尤其当水印边缘压在人物皮肤或者布料纹理上时。实现时可以选择预训练VGG的 relu1_2、relu2_2、relu3_3 三层特征做距离约束，让模型在语义特征层面也贴近原图。

import torch import torch.nn.functional as F def hybrid_loss(pred, target, mask=None): """ pred: 模型输出的去水印图 [B,3,H,W] target: 干净原图 [B,3,H,W] mask: 水印区域mask [B,1,H,W]，可空 """ # L1 保边缘，L2 促平滑，默认各0.5权重 loss_l1 = F.l1_loss(pred, target) loss_l2 = F.mse_loss(pred, target) loss = 0.5 * loss_l1 + 0.5 * loss_l2 # 如果有水印mask，对mask区域额外加权，让模型把火力集中在水印上 if mask is not None: region_loss = F.l1_loss(pred * mask, target * mask) loss = loss + 0.3 * region_loss return loss

这个 mask 加权的小动作很值得用。不加权时模型需要自己去“发现”水印区域，加了之后相当于告诉它“重点在这儿”，训练速度明显加快，也能减少水印区域反复残留的问题。要注意的是 mask 加权系数不要超过 0.5，否则非水印区域的背景纹理失去约束，生成结果会出现不自然的“油画感”。

4.2 学习率调度和权重命名里的信息

优化器用AdamW，初始学习率按 batch size 缩放，常见做法是 lr=1e-4 起步，配合批大小 8 到 16。学习率调度我建议用 warmup 加余弦退火：前 5 个 epoch 线性升到 1e-4，之后按余弦曲线降到 1e-6。为什么不直接固定学习率？因为训练后期模型已经在微调纹理细节，学习率太大时L1和L2的梯度会把刚恢复出来的边缘再次抹掉。这一点在去水印任务上尤其明显，我自己踩过一次后才老实加上退火。

权重文件 v7ms2_ep49.pd 这个命名里其实藏着信息：v7ms2 是训练版本号，ep49 代表第49轮保存的权重。这就意味着方案训练过程至少到49轮才进入稳定收敛状态，不是你跑20轮就能复现效果的。训练轮次这种参数，常规设置是60到80轮，配合余弦退火让后半段慢慢收敛。

还有一个容易被忽略的工程细节是EMA。给模型参数做指数滑动平均，能明显提升最终推理时的重建稳定度。代码写起来就几行：

class EMAHelper: def __init__(self, model, decay=0.999): self.decay = decay self.shadow = {k: v.detach().clone() for k, v in model.state_dict().items()} def update(self, model): for k, v in model.state_dict().items(): # shadow慢慢朝当前参数靠拢，相当于给参数做平滑 self.shadow[k] = self.decay * self.shadow[k] + (1 - self.decay) * v.detach() def apply(self, model): model.load_state_dict(self.shadow, strict=False)

EMA 的 decay 值设到 0.999 以上比较合适，小于 0.99 时平滑作用几乎起不到。推理时把 EMA 参数替换进去，验证集 PSNR 一般能稳定涨 0.1～0.3。这个提升幅度不算大，但比赛里名次经常就卡在这个量级上。

5. 训练去水印模型避坑：三条翻车记录和一条显存教训

5.1 水印去掉了，整张图却发灰

现象是模型输出后，水印区域确实干净了，但整张图的对比度变低，暗部发灰、亮部发闷，人眼看着像隔了层雾。这个现象在验证集PSNR上反而不容易发现，因为PSNR只算像素差，灰蒙蒙的图像素差未必大。

原因出在L2 loss占比过高，加上增强管线里光度扰动太激进，模型为了最小化平均误差，选择输出一个“最保守”的中间灰度。这本质是回归任务的均值收缩问题。解决方式是我在前面提到的 hybrid loss 结构调整：把L2比重降到 0.3～0.5，同时加入感知 loss，让模型在语义特征层面和原图对齐。另外把颜色类增强的幅度调低，亮度变化范围从 ±0.3 收窄到 ±0.15。

5.2 浅色水印残留，细看还有一层“白影子”

现象是强水印去得很干净，但透明度在 0.2 以下的浅水印区域会留下淡淡的白边，放大看能看出字迹轮廓。原因有两个层面：一是合成训练数据时 alpha_range 最小值设到了 0.3 以上，模型没怎么见过低透明度水印；二是注意力模块把权重大量分配给了明显的大水印，对弱水印响应不足。

解决方式分两步。第一步把合成数据 alpha_range 下限降到 0.1，但此时要同时提高 mask 加权系数，让模型知道这个淡淡的区域也要去处理。第二步是在训练后期把低透明度水印样本的出现比例提高，相当于做一个简单的课程学习：前期主要去重水印，后期精细化处理轻水印。我一般会把两者比例控制在 7:3，重水印为主，轻水印为辅助。

5.3 验证集loss不降反升，训练集却一直在降

这个现象就是典型的过拟合。增强强度不够是主要嫌疑：如果训练数据里水印都是同一字号、同一透明度、同一位置，模型背都能背下来，验证集换一版新水印后立刻打回原形。解决方式是加大几何扰动和光度扰动，把随机裁剪尺度范围扩大，让水印的位置分布更散。另外把数据集里的字体库补齐，至少包含衬线、无衬线和手写三类风格，否则模型会对特定笔画结构过拟合。

调完增强之后还要注意验证集评估方式：验证时不要用训练时同一批合成参数生成的水印，单独用另一批字体和透明度范围来测，这样才能真实反映泛化能力。用同一批生成参数做验证，PSNR再高都有自欺欺人的成分。

5.4 一次OOM崩溃暴露出的显存管理问题

现象是训练到第30轮左右，batch size 为 16 时直接 Out of Memory，进程崩溃。原因比想象中朴素：混合注意力模块里空间注意力要对 512x512 的特征图做 7x7 卷积，这个分支的中间变量占用了大量显存，而且到了训练后期模型EMA开启后，内存开销又增加了一份。

解决方式有几种可选组合。一是把 batch size 从 16 降到 8；二是对空间注意力特征图做一次降采样再上采样，把 7x7 卷积作用在 256x256 分辨率上；三是开启 Paddle 的梯度累积功能，用小 batch 模拟大 batch。方案本身没有开大batch，这恰恰说明显存控制是靠模块设计而不是硬件堆出来的，低显存环境跑这个方案完全可行，2300 元左右的消费级显卡就能覆盖。

6. 复现冠军效果：一个低成本验证技巧和它的意义

整个项目拆到最后，有一个细节我觉得比任何模块都值得分享：如何在没有原图的情况下验证去水印模型效果。比赛有标准测试集，但实际项目中你会遇到大量“只有带水印图、没有干净原图”的场景。我养成的习惯是单独攒一个“自拍盲测集”：找 20 到 30 张不同来源的带水印照片，包含白字、黑字、彩色logo、贴纸四类，跑完推理后不计算任何指标，直接拼成大图在屏幕上放大看边缘。

这个习惯救过我一次。有一次模型验证集PSNR到了 33.5，我拿盲测集一看，发现浅灰色水印的上缘残留了一条半像素宽的白线，缩略图看不出，放大到 200% 后清晰可见。PSNR 对这种半像素级别的系统性误差基本无感。从那以后我每次训练结束都强制走一遍盲测流程，先看边缘再看整体，最后才信指标。如果你也想快速跑通整个推理链路，思路很简单：加载 v7ms2_ep49.pd 权重，输入归一化到 0～1，前处理统一缩放到 512x512，推理后再缩回原尺寸。跑完第一版、确认图像没有发灰发糊，再回头调loss和增强参数才有意义。这套方案包的代码和权重都是现成的，适合想快速复现、以及想在基础上替换模块做消融实验的人直接动手。希望帮到你。

本文还有配套的精品资源，点击获取