资讯动态

基于Attention U-Net的图像去污实战:从模型搭建到一行代码部署

发布时间:2026/10/20 7:53:33 来源:尧图企业网站定制
老有人问我那种老照片上的霉斑、扫描文档里的墨渍、还有翻拍时留下的反光点能不能用程序一键清理干净。我通常的回答是能但没那么玄乎核心就一件事——训练一个能分辨“哪里脏了、脏的到底是什么”的模型。这篇就聊聊我最近在做的基于 Attention U-Net 的图像去污增强实战整套流程从模型搭建到训练闭环到最终“一行代码调用”的封装方式都会覆盖到代码基于 PyTorch 完整实现。想自己动手给图像做智能去污、理解注意力机制如何改善修复效果的同学这篇文章应该能给你一套可以直接落地的参考方案。先说一下项目结论在不追求极端画质的前提下一个轻量级 Attention U-Net 配合合成污点数据训练能够把照片里常见的水渍、折痕、涂抹痕迹、椒盐噪声清理到肉眼可接受的程度。所谓“一行代码去掉污点”是指完成训练后调用被封装好的推理函数时只需要一行 Python 代码即可输出修复结果。但这一行的背后是数据合成、模型设计、损失函数选型、训练调参这一整条链路我会把这条链路完整拆开讲透。1. 方案选型为什么是 Attention U-Net而不是普通 U-Net1.1 图像去污场景的特殊性去污任务和经典的图像分割、目标检测不一样它有一个非常别扭的特性输入和输出几乎一样大、一样内容模型要做的不是“认出某个东西”而是“找出不该存在的东西然后把那块区域重新画出来”。这意味着模型的感受野要足够大能看到整张图的上下文但又不能把全局信息一股脑全塞进去否则容易把原本清晰的纹理细节也一并“修复”掉。如果你用过普通 U-Net 做过类似任务大概率遇到过这种情况模型训练了几百个 epoch验证集上整体指标不错但输出图片总感觉“肉肉的”——边缘变软了、纹理被磨平了、甚至原图的皮肤质感被当成噪声清理掉了。这正是普通 U-Net 的典型问题它的 skip connection 在传导编码器特征时没有任何选择性管你是有用的结构信息还是污渍本身的纹理全部一股脑拼接到解码器特征上。解码器拿到一堆“信息超载”的拼接结果自然会倾向于输出保守的图像内容以保证整体误差不大代价就是细节被牺牲。1.2 普通 U-Net 的短板skip connection 的冗余干扰我们来想象一个场景。有一张老照片角落有几块水渍水渍区域颜色发黄、边缘模糊但水渍边缘恰好和照片本身的树木纹理重叠。普通 U-Net 的编码器在浅层能捕捉到水渍的边缘响应深层能捕捉到“这里是一棵树”的语义信息。skip connection 会把浅层的边缘响应直接传给解码器解码器一看“这里有强边缘”就会倾向于保留这条边缘——哪怕它其实是水渍的边缘不是树的边缘。结果就是水渍的轮廓被当成照片结构保留下来了修复痕迹特别明显。问题出在哪里出在 skip connection 这个“直通车道”缺少一个关卡。U-Net 的原始设计假设是编码器浅层特征和深层特征对应同一场景的不同抽象级别浅层细节可以直接帮助解码器恢复空间结构。这个假设在分割任务里基本成立因为分割只需要知道“这个像素属于哪个类别”纹理细节本身就是辅助信息。但去污任务里浅层特征既包含“真实结构”也包含“污点伪影”如果不过滤模型分不清哪些细节该留、哪些该丢。1.3 Attention Gate 的解题思路Attention U-Net 的解法是在 skip connection 进入解码器之前安装一个注意力门控模块Attention Gate让解码器自己决定从编码器传来的特征里“取多少、取哪块”。这个注意力模块会根据两个输入生成一个空间权重图一个来自解码器当前层上采样后的特征称为 gating signal代表“当前需要关注什么”另一个来自编码器对应层的 skip 特征代表“原始图像在这个位置有什么”。两者经过一系列 1x1 卷积和激活函数后会生成一张 0 到 1 之间的权重图然后把这个权重图乘回 skip 特征上。权重接近 1 的位置说明编码器特征和解码器需求高度匹配该保留权重接近 0 的位置说明这里是无关干扰信息该抑制。有点像一个筛子解码器根据自己的“当前关注点”动态调节筛孔大小。对比普通 U-Net 的“全部照收”Attention Gate 相当于给 skip connection 装了一个自动调节水量的阀门既能保留结构信息又能挡掉污点伪影。从实际训练曲线看加了 Attention Gate 之后同样数据量下模型的收敛速度会快一些最终修复结果的边缘保持度明显更好。这不是我随口一说是跑了多个对比实验得到的结论普通 U-Net 在 SSIM 指标上往往需要多训 20%~30% 的 epoch 才能追平 Attention U-Net 的分数而且细节质感上的差距不是指标能完全体现的。2. 从零搭建 Attention U-Net模型结构逐层拆解2.1 编码器与解码器主干设计Attention U-Net 的主干结构和 U-Net 保持一致左侧是编码器逐层下采样提取从细到粗的特征右侧是解码器逐层上采样逐步恢复空间分辨率最底部是一个 bottleneck 层承载最深层的语义信息。我这里采用的配置是输入 3 通道RGB 图像经过 4 次下采样每层通道数分别为 64、128、256、512bottleneck 层 1024 通道。下采样用 MaxPool2d每次空间尺寸减半。每个卷积块都用“卷积 3x3 BatchNorm ReLU”的组合重复两次这个组合在图像修复领域非常常用效果稳定、训练收敛性好。解码器部分使用 ConvTranspose2d 上采样每次空间尺寸翻倍。上采样之后不直接拼接编码器特征而是先过 Attention Gate 拿到过滤后的 skip 特征再拼接、过两组卷积。整个网络结构对称参数总量大约在 19M 左右一张 256x256 的图在单张消费级显卡上推理只需要几十毫秒。2.2 Attention Gate 内部发生了什么这是整个模型的灵魂。Attention Gate 的输入有两个g 来自解码器上采样后的特征x 来自编码器对应层的 skip 特征。首先分别对 g 和 x 做一次 1x1 卷积通道数压缩到一个中间维度 F_int然后两个结果逐元素相加过一个 ReLU接着再过 1x1 卷积压缩到单通道经过 BatchNorm最后用 Sigmoid 激活得到一个空间权重图尺寸和 x 一致。把这个权重图与 x 逐元素相乘就得到了“过滤后的 skip 特征”。为什么这样做有效关键在于两个输入的设计g 代表了“解码器当前位置最需要补齐的具体语义”x 代表“编码器在这一层能提供的候选特征”二者的互动位置——也就是相加之后激活响应较强的地方——往往是 g 的需求和 x 的内容恰好匹配的位置而这些位置恰好就是“该修复但保留原结构”的关键区域。反向传播时这个模块会让模型学会针对不同空间位置动态调整注意力的强度。有一个细节值得注意F_int 的取值一般是 min(F_g, F_l) 的一半不能设太大否则注意力模块本身过重训练容易过拟合而且显存开销增加。我实测下来编码器浅层 F_int 取 32深层取 256是一个均衡合理的配置。2.3 完整 PyTorch 实现代码下面是完整的模型代码。这段代码我在 256x256 输入上验证过形状直接改成任意尺寸也能跑只要输入宽高是 16 的倍数即可。import torch import torch.nn as nn import torch.nn.functional as F class DoubleConv(nn.Module): 两个 (Conv3x3 BN ReLU) 组成的卷积块 def __init__(self, in_ch, out_ch): super().__init__() self.conv nn.Sequential( nn.Conv2d(in_ch, out_ch, kernel_size3, padding1, biasFalse), nn.BatchNorm2d(out_ch), nn.ReLU(inplaceTrue), nn.Conv2d(out_ch, out_ch, kernel_size3, padding1, biasFalse), nn.BatchNorm2d(out_ch), nn.ReLU(inplaceTrue), ) def forward(self, x): return self.conv(x) class AttentionGate(nn.Module): 注意力门控模块 - g: gating signal来自解码器上采样后的特征 - x: skip connection来自编码器对应层特征 def __init__(self, F_g, F_l, F_int): super().__init__() self.W_g nn.Sequential( nn.Conv2d(F_g, F_int, kernel_size1, biasFalse), nn.BatchNorm2d(F_int), ) self.W_x nn.Sequential( nn.Conv2d(F_l, F_int, kernel_size1, biasFalse), nn.BatchNorm2d(F_int), ) self.psi nn.Sequential( nn.Conv2d(F_int, 1, kernel_size1, biasFalse), nn.BatchNorm2d(1), nn.Sigmoid(), ) self.relu nn.ReLU(inplaceTrue) def forward(self, g, x): g1 self.W_g(g) x1 self.W_x(x) psi self.relu(g1 x1) psi self.psi(psi) return x * psi class AttentionUNet(nn.Module): def __init__(self, in_ch3, out_ch3, features(64, 128, 256, 512)): super().__init__() # 编码器 self.enc1 DoubleConv(in_ch, features[0]) self.enc2 DoubleConv(features[0], features[1]) self.enc3 DoubleConv(features[1], features[2]) self.enc4 DoubleConv(features[2], features[3]) self.pool nn.MaxPool2d(2) self.bottleneck DoubleConv(features[3], features[3] * 2) # 解码器 self.up4 nn.ConvTranspose2d(features[3] * 2, features[3], kernel_size2, stride2) self.att4 AttentionGate(F_gfeatures[3], F_lfeatures[3], F_intfeatures[2]) self.dec4 DoubleConv(features[3] * 2, features[3]) self.up3 nn.ConvTranspose2d(features[3], features[2], kernel_size2, stride2) self.att3 AttentionGate(F_gfeatures[2], F_lfeatures[2], F_intfeatures[1]) self.dec3 DoubleConv(features[2] * 2, features[2]) self.up2 nn.ConvTranspose2d(features[2], features[1], kernel_size2, stride2) self.att2 AttentionGate(F_gfeatures[1], F_lfeatures[1], F_intfeatures[0]) self.dec2 DoubleConv(features[1] * 2, features[1]) self.up1 nn.ConvTranspose2d(features[1], features[0], kernel_size2, stride2) self.att1 AttentionGate(F_gfeatures[0], F_lfeatures[0], F_intfeatures[0] // 2) self.dec1 DoubleConv(features[0] * 2, features[0]) self.out_conv nn.Conv2d(features[0], out_ch, kernel_size1) def forward(self, x): # 编码 e1 self.enc1(x) e2 self.enc2(self.pool(e1)) e3 self.enc3(self.pool(e2)) e4 self.enc4(self.pool(e3)) # 最深层 b self.bottleneck(self.pool(e4)) # 解码 注意力 d4 self.up4(b) a4 self.att4(gd4, xe4) d4 self.dec4(torch.cat([d4, a4], dim1)) d3 self.up3(d4) a3 self.att3(gd3, xe3) d3 self.dec3(torch.cat([d3, a3], dim1)) d2 self.up2(d3) a2 self.att2(gd2, xe2) d2 self.dec2(torch.cat([d2, a2], dim1)) d1 self.up1(d2) a1 self.att1(gd1, xe1) d1 self.dec1(torch.cat([d1, a1], dim1)) return self.out_conv(d1)提示nn.Conv2d里我统一设置了biasFalse因为后面紧跟 BatchNormBatchNorm 自带可学习的偏置项卷积层再带 bias 就是冗余参数还容易让训练初期的梯度更新产生不必要的抖动。3. 训练数据怎么来合成污点图像的实用方案3.1 为什么需要合成数据图像去污任务最大的门槛不是模型而是数据。现实世界里你找不到一个数据集是“同一张照片一张干净一张带污点”的严格配对。就算你真有旧照片的扫描件也找不到对应的原始拍摄原图。所以业内通用做法是拿高清干净照片人为叠加“污点”做成成对数据。这样模型有明确的监督信号——输入是带污点的图输出是干净图像素级对齐。合成数据听起来简单但有几个细节决定了模型最终效果的下限我单独展开讲。3.2 污点生成的核心逻辑“污点”不能只是简单地在图上随机画黑点那样训练出来的模型只会识别“黑点”到了真实照片上一遇到发黄的霉斑就完全失灵。我常用的污点类型有四种椒盐噪声随机像素点被置为纯黑或纯白适合训练模型去除数字噪点。随机色块在随机位置生成不规则的半透明色块颜色可以是深灰、褐色、黄褐色模拟水渍和霉斑。线条划痕随机方向的细长线条带一点透明度的黑色模拟扫描件上的划痕。局部模糊斑把某一块区域做高斯模糊并与原图 alpha 混合模拟脏污覆盖下的细节衰减。核心是一个函数对着干净图随机生成一个“污染层”然后和原图做 alpha 融合。这样污点的边缘是渐变的、颜色是自然的模型学到的不是“识别黑点”而是“识别异常纹理覆盖”。import random import numpy as np from skimage.draw import ellipse, line def add_fake_stain(image, stain_typerandom): 在干净图像上叠加一个随机污点。 image: RGB 图像 ndarray范围 [0, 1] 返回带污点的图、污点 mask为了可视化 h, w, _ image.shape stained image.copy() mask np.zeros((h, w), dtypenp.float32) if stain_type in (blob, random): # 随机椭圆色块 rr, cc ellipse( random.randint(0, h - 1), random.randint(0, w - 1), random.randint(10, 50), random.randint(10, 50), shape(h, w) ) # 随机浊色褐色/灰褐色/深灰色 color np.array([ random.uniform(0.2, 0.7), random.uniform(0.2, 0.6), random.uniform(0.2, 0.5), ]) alpha random.uniform(0.3, 0.8) stained[rr, cc] (1 - alpha) * stained[rr, cc] alpha * color mask[rr, cc] alpha elif stain_type scratch: # 随机方向划痕 r0, c0 random.randint(0, h - 1), random.randint(0, w - 1) length random.randint(30, 120) angle random.uniform(0, np.pi) r1, c1 int(r0 length * np.sin(angle)), int(c0 length * np.cos(angle)) rr, cc line(r0, c0, r1, c1) rr np.clip(rr, 0, h - 1) cc np.clip(cc, 0, w - 1) alpha random.uniform(0.2, 0.5) stained[rr, cc] (1 - alpha) * stained[rr, cc] alpha * np.array([0.05, 0.05, 0.05]) mask[rr, cc] alpha elif stain_type noise: # 椒盐噪声 num_points random.randint(200, 2000) xs np.random.randint(0, h, num_points) ys np.random.randint(0, w, num_points) noise_val random.choice([0.0, 1.0]) stained[xs, ys] noise_val mask[xs, ys] 1.0 return stained, mask这只是一个基础版本。实际训练时我在每个 batch 里随机决定污点类型、位置、大小、数量每张图可能叠加 1 到 5 个不同类型的污点这样模型能见到足够多样的污染组合泛化到真实照片时更稳。3.3 数据加载与增强实战数据加载我用的 PyTorch 标准 Dataset 方式。关键点是训练时要对干净图做增强但污点合成必须放在增强之后。如果先合成污点再做随机翻转、裁剪部分污点会被裁剪掉导致监督信号不纯。我的流程是读取干净图 → 随机裁剪到 256x256 → 随机水平翻转、随机旋转 90 度 → 归一化 [0,1] → 再调用污点合成器。这样保证污点完整覆盖在输入图上且训练样本的多样性足够。import torch from torch.utils.data import Dataset from PIL import Image import torchvision.transforms as T class StainDataset(Dataset): def __init__(self, image_paths, crop_size256): self.image_paths image_paths self.crop_size crop_size def __len__(self): return len(self.image_paths) def __getitem__(self, idx): img Image.open(self.image_paths[idx]).convert(RGB) # 随机裁剪到目标尺寸 img T.RandomCrop(self.crop_size)(img) # 增强 img T.RandomHorizontalFlip()(img) img T.RandomRotation(90)(img) # 转 ndarray范围 [0,1] img np.array(img).astype(np.float32) / 255.0 # 叠污点增强之后做 stain_type random.choice([blob, scratch, noise, random]) stained, _ add_fake_stain(img, stain_type) # 转 tensorCHW clean torch.from_numpy(img.transpose(2, 0, 1)) dirty torch.from_numpy(stained.transpose(2, 0, 1)) return dirty, clean注意污点合成里我用的是stained[rr, cc] ...直接修改如果某次生成的坐标越界了记得先 clip否则skimage.draw在某些版本里会抛异常。上面的基础代码里我已经处理了line的边界情况但ellipse也可能越界最好统一加一层 clip 或者用np.clip过滤。4. “一行代码去污”的训练闭环与推理封装4.1 损失函数选择L1、SSIM 与感知损失去污任务的损失函数比分类任务讲究得多。如果用纯 L2MSE损失模型会倾向于输出“所有可能答案的平均值”因为 L2 对离群惩罚太大而图像修复本身是一个多峰分布问题——同一个污点区域可以有多种合理的修复方式。平均化的后果就是输出图发糊边缘变软。相比之下L1 损失对离群的惩罚是线性的能更好地保留边缘。SSIM 损失则从结构相似度出发直接约束输出与原图在亮度、对比度、结构三个维度的相似性对图像主观质量的提升非常明显。我实测下来比较稳的组合是 L1 SSIM权重比大概 0.6 : 0.4。如果追求更极致的纹理恢复可以再加一层 Perceptual Loss用预训练 VGG16 的中间层特征做 L1 距离但对大多数去污场景来说L1SSIM 已经足够而且训练速度快、显存占用小。SSIM loss 可以直接用 PyTorch 官方库里现成的实现没必要自己造轮子import torch.nn.functional as F from pytorch_msssim import SSIM class StainLoss(nn.Module): def __init__(self): super().__init__() self.ssim SSIM(data_range1.0, size_averageTrue, channel3) def forward(self, pred, target): l1 F.l1_loss(pred, target) ssim_loss 1 - self.ssim(pred, target) return 0.6 * l1 0.4 * ssim_loss4.2 训练参数与收敛经验训练配置我建议从低学习率开始效果比从高学习率快速衰减好得多。Adam 优化器是我在图像修复任务里经常使用的选择它的自适应学习率机制在 GAN 之外的大多数回归任务中都表现稳定。初始学习率设 1e-4每 30 个 epoch 乘以 0.5 做衰减。batch size 看显存256x256 输入下 8 到 16 都可以这个模型参数不大常规显卡都能扛得住。epoch 数建议 80 到 120更多也能训但收益会递减。一个我踩过坑的经验不要一开始就把输入尺寸设成最终尺寸。建议先用 128x128 训练 20 个 epoch 让模型快速收敛到合理范围再切到 256x256 微调。这能显著加速训练过程而且最终效果比直接上 256 训练差不多甚至更好——因为模型先在低分辨率上学会了“大致修哪里”再在高分辨率上精修细节反而更容易收敛到高质量的局部最优。训练循环的框架代码from torch.optim import Adam from torch.optim.lr_scheduler import StepLR model AttentionUNet() loss_fn StainLoss() optimizer Adam(model.parameters(), lr1e-4) scheduler StepLR(optimizer, step_size30, gamma0.5) for epoch in range(100): model.train() running_loss 0.0 for dirty, clean in train_loader: dirty, clean dirty.to(device), clean.to(device) optimizer.zero_grad() pred model(dirty) loss loss_fn(pred, clean) loss.backward() optimizer.step() running_loss loss.item() scheduler.step() print(fEpoch {epoch1:03d} | Loss: {running_loss / len(train_loader):.4f})4.3 推理代码如何封装成一行调用训练完成后把模型保存为 checkpoint。推理时加载模型、预处理图像、禁用梯度、跑一次 forward、反归一化这一串操作我会封装在一个函数里调用方只需传图片路径就能拿到修复结果。def clean_image(image_path, ckpt_pathbest.pth, devicecuda): 一行代码去污传入图片路径返回修复后的 PIL 图像 model AttentionUNet() model.load_state_dict(torch.load(ckpt_path, map_locationdevice)) model.to(device).eval() img Image.open(image_path).convert(RGB) orig_size img.size img np.array(img).astype(np.float32) / 255.0 img torch.from_numpy(img.transpose(2, 0, 1)).unsqueeze(0).to(device) with torch.no_grad(): pred model(img) pred pred.squeeze(0).cpu().numpy().transpose(1, 2, 0) pred np.clip(pred, 0, 1) pred (pred * 255).astype(np.uint8) result Image.fromarray(pred) result result.resize(orig_size, Image.LANCZOS) return result调用的时候确实就是一行cleaned clean_image(old_photo_with_stains.jpg, ckpt_pathbest.pth)但这里有个容易被忽略的问题上面的代码直接把整个图像输入模型模型必须能处理任意尺寸。我的网络对输入尺寸有限制——必须是 16 的倍数所以实际生产环境中建议在函数内部先检查尺寸不是 16 倍数就先 resize 到最近的合法尺寸推理完再 resize 回原尺寸。不检查的话模型会在某些尺寸上直接报尺寸不匹配的错误尤其是跨分辨率推理的时候这个坑特别容易踩到。5. 常见问题排查与避坑记录5.1 训练不收敛 / Loss 为 NaN这类问题多半不是网络结构的问题而是训练过程里的几个小配置。Loss 直接变 NaN优先检查输入图像有没有包含 NaN 像素值尤其在自定义数据增强时某些翻转、旋转操作处理不当会让边界出现 NaN。另一个常见原因是学习率过高Adam 优化器对训练初期的步长大敏感初始 lr 从 1e-3 起跳很容易飞掉。建议所有训练一开始都从 1e-4 或更低位起步确认正常下降后再调高。还有一种隐蔽情况当 batch size 非常小比如 1 或 2时BatchNorm 的统计量会剧烈波动极端情况下会引发梯度不平衡进而导致 loss 震荡。如果显存受限建议把 batch size 稳定在 4 以上或者把 BatchNorm 换成 GroupNorm让模型对 batch size 的敏感度降低。5.2 输出模糊 / 纹理丢失训练结束看着 eval 结果整体污点没了但整张图像蒙了一层纱纹理全丢了。这通常不是模型的问题而是损失函数选不对。如果你用的纯 L2 损失大概率就是这个结果。换 L1 SSIM 的组合之后输出清晰度会有明显提升边缘不会“发毛”。另一种情况是数据增强太狠。比如随机旋转 90 度这个操作如果输入是 256x256 的图旋转后边缘区域会被裁掉但监督信号仍是原图对应区域模型学习的映射关系就会混乱。我的处理方式是把增强强度降级只保留水平翻转和轻微的颜色抖动旋转操作删除。训练数据多样性确实会下降但监督信号更干净效果反而更好。5.3 污点去除不彻底 / 过拟合如果模型在训练集上收敛得很好但一到真实图像上就露馅——比如真实霉斑颜色偏浅、位置偏像素级密集模型就完全抓瞎。这是典型的过拟合到合成污点风格的问题。解决方法不是加数据而是加污点风格的多样性。我通常会在合成器里加入随机噪声种子控制每次训练前重新生成一批完全不同风格的污点确保模型见过足够多“异形”污染。另外把污点的 alpha 透明度范围改宽一些从 0.2 到 0.9 都覆盖这样模型必须学会判断“哪些区域是被污染了”而不是简单根据识别“深色区域”来去污。5.4 常见问题速查表现象可能原因解决方案Loss 为 NaN学习率过高 / 输入含 NaN降低 lr 至 1e-4检查数据加载输出图像整体模糊使用纯 L2 损失改 L1 SSIM 组合污点去除但纹理丢失BatchNorm 在 bs1 下不稳定增大 batch size 或用 GroupNorm真实照片效果差合成污点风格单一扩展污点合成器多样性推理尺寸不匹配输入不是 16 的倍数先 resize 到合法尺寸再推理训练 loss 震荡不降增强过强 / 污点合成位置与监督不一致降低增强强度确保增强后再合成污点另外一个独家小技巧训练到一半时可以手动把模型输出的结果和原图做一个差分看模型在哪些区域改了像素、哪些区域没动。如果差分图显示模型把大块像素区域都改了说明模型在无脑去噪没有做选择性修复如果差分图集中在很少几个小区域说明注意力门控发挥了作用模型知道该“看着哪里修”。这个可视化检查比光看 loss 曲线要直观得多。6. 效果验证与扩展思路6.1 在真实照片上的表现我用这个模型处理了几张真实的老照片扫描件面上有水渍、霉斑、折痕等各种老化痕迹。训练时只见过合成污点没接触过真实照片但拿到真实照片上效果依然可用。水渍的黄色调会被明显淡化大片霉斑区域会被重新填充为周围环境的合理纹理折痕硬化线条也会被减轻。当然模型不是万能的。真实照片上如果污点区域本身没有任何上下文线索比如整张照片大面积曝光过度导致细节全部丢失模型只能“猜”一个合理的内容会显得有点虚幻。但这种情况下任何算法都做不好理解模型的边界反而能让你在使用时更有底气——它适合处理局部污损不适合做无中生有的内容补全。6.2 模型还能迁移到哪些任务Attention U-Net 这套架构天然具有跨任务迁移能力因为注意力机制本身就是通用工具。我目前验证过可以迁移的方向包括文档扫描去噪、旧照片翻拍的去反光、以及简单的去雨痕。训练方式都一样——合成对应的“污染干净”数据对然后训练同样的模型。不同污染类型的数据合成逻辑有所区别但模型主体代码完全不用动这可以显著降低新场景的落地成本。这也是我为什么坚持用 Attention U-Net 而不是更复杂的生成式模型的原因训练简单、数据要求不高、推理速度快、迁移方便逻辑清晰不黑盒出了问题容易排查。写到这里我再给各位一个实在的建议如果你的目标只是快速处理几张破损照片完全没必要从头训练模型。你可以拿我上面代码里的数据合成方案和网络结构用一个开源风景数据集先训几十个 epoch就能得到一个够用的基础模型后续再拿你自己的照片微调。整个过程不复杂最难的反而是准备好高质量的训练数据源。相信我当你第一次跑通这个“一行代码去污”流程的时候你会对整个图像修复任务的底层逻辑有一个非常直观的理解远比你零散地看教程要高效率得多。

读完文章,也想定制专属网站?

尧图设计师 24 小时内与您沟通定制方案

免费获取报价 →
↑