1. 从“输入和标签都是图片”说起去雨任务的数据集到底特殊在哪图像去雨这个方向很多人第一次接触时都会卡在一个很朴素的问题上数据集到底长什么样如果你之前做过分类任务习惯的是“一张图 一个类别标签”做过分割任务习惯的是“一张图 一张掩码图”。而去雨任务属于图像到图像的映射输入是一张带雨痕的图标签是同一场景下没有雨的那张干净图。听起来简单但真正动手搭 Dataset 的时候坑远比想象中多。我最初拿到一个去雨数据集时第一反应是“不就是两个文件夹一个放带雨图一个放干净图按文件名配对嘛”。结果打开目录一看带雨图有 1200 张干净图只有 400 张文件名还完全对不上。这就是去雨数据集最典型的特征之一它不是严格的一一对应关系而是“一张干净图对应多张不同雨痕的合成图”。这个结构直接决定了你在写 PyTorch 的Dataset类时配对逻辑不能想当然。为什么会出现这种结构因为真实世界里“同一场景、同一角度、同一光照只有雨不一样”的成对数据几乎拍不到。雨是动态的光线在变相机稍微抖一下场景就变了。所以主流做法是先采集干净背景图然后用物理渲染或者图层叠加的方式人工合成不同密度、不同方向、不同雨滴大小的雨痕生成多张带雨图。这样一来一张干净图就派生出了好几张训练样本。理解这一点非常关键因为它影响三件事第一你的数据配对逻辑必须支持“一对多”第二划分训练集和验证集时必须按“干净图”来划分而不是按“带雨图”划分否则同一张干净图派生出的带雨图会同时出现在训练和验证里造成数据泄漏第三评估指标的可信度依赖于验证集里干净图的多样性而不是带雨图的数量。我见过不少新手在这里翻车验证集 loss 低得离谱以为模型学得很好结果一放到真实雨天照片上就完全失效。排查半天才发现训练集和验证集里存在同一场景的不同雨痕版本模型其实是在“背场景”而不是在“学去雨”。这个坑非常隐蔽因为从带雨图的文件名上根本看不出来它们共享同一张干净背景。所以这一节我想先把认知层面的事情讲透去雨数据集的核心不是“图片数量”而是“场景数量”和“雨痕多样性”的平衡。你后面写 Dataset、做增强、划分数据集所有决策都要围绕这个认知展开。下面我会从目录结构设计、配对逻辑、数据划分三个角度把这件事拆开讲。1.1 常见的两种目录组织方式与各自的坑去雨数据集在磁盘上的组织方式我见过的主流有两种。第一种是双文件夹平行结构dataset/ rainy/ scene01_rain_a.png scene01_rain_b.png scene02_rain_a.png ... clean/ scene01.png scene02.png ...这种结构的特点是干净图只有一份带雨图有多份通过文件名前缀比如scene01来关联。写 Dataset 时你需要从带雨图文件名里解析出场景 ID再去 clean 目录里找对应的干净图。优点是磁盘占用小干净图不重复存储缺点是配对逻辑依赖命名规范一旦命名不统一就会出错。第二种是成对文件夹结构dataset/ train/ rainy/ scene01_rain_a.png scene01_rain_b.png clean/ scene01_rain_a.png scene01_rain_b.png注意这里 clean 目录里其实是同一张图复制了两份只是文件名和带雨图一一对应。这种结构写 Dataset 最简单直接按文件名配对即可但磁盘占用会成倍增加。如果你的干净图是 2K 分辨率、几千张复制一份可能就多占几个 G。我个人的选择是训练阶段用第一种配合内存缓存如果磁盘实在紧张第二种只在小规模调试时用。原因很简单去雨模型的训练轮次通常很多磁盘 IO 会成为瓶颈与其复制文件不如在 Dataset 里做一次映射缓存。这里有个细节要注意解析场景 ID 时不要用简单的字符串切割。比如scene01_rain_a.png和scene01_rain_b.png如果你用split(_)[0]拿到scene01看起来没问题但如果有的数据集命名是scene_01_rain.png切割结果就错了。稳妥的做法是用正则表达式提取“场景标识 雨痕标识”两部分并且把映射关系在初始化时一次性建好存成一个字典避免每个__getitem__都去遍历目录。1.2 数据划分为什么必须按干净图切分前面提到按干净图划分这里展开说清楚操作方式。假设你有 400 张干净图每张派生出 3 张带雨图总共 1200 张带雨图。正确的划分方式是先把 400 张干净图按 8:1:1 切成训练 320、验证 40、测试 40然后每张干净图对应的 3 张带雨图跟着它一起进对应的集合。这样训练集有 960 张带雨图验证集和测试集各 120 张。错误的方式是直接把 1200 张带雨图打乱按比例切分。这样同一张干净图的 3 张带雨图很可能被分到不同集合验证集里的场景在训练集里出现过指标虚高。实现上我通常会在 Dataset 初始化之前先构建一个scene_id - [rainy_paths]的字典然后对scene_id列表做划分最后展开成实际的图片路径列表。这样逻辑清晰也不容易出错。划分时记得固定随机种子保证每次实验的划分一致否则不同轮次的对比就没有意义了。提示如果你的数据集本身没有提供场景 ID只能从文件名推断那一定要先写个小脚本统计一下“干净图数量”和“带雨图数量”的比例以及文件名前缀的分布情况。我遇到过前缀解析错误导致 1200 张图只匹配上 300 张的情况训练时 loss 正常下降但模型效果极差排查了很久才发现是配对漏了。2. 写一个能跑通的去雨 Dataset从__init__到__getitem__认知理清之后进入代码层面。PyTorch 的Dataset类看起来简单就三个方法但去雨任务的特殊性让每个方法都有讲究。我下面给出的实现是我自己在多个项目里迭代过的版本兼顾了灵活性、性能和可读性。先明确目标输入是带雨图路径输出是(rainy_tensor, clean_tensor)这样一对张量形状都是[3, H, W]数值范围归一化到[0, 1]或者标准化到[-1, 1]具体取决于你后面用的损失函数和预训练模型。2.1__init__里该做什么、不该做什么__init__的职责是建立索引和映射而不是加载图片数据。这一点很多人会搞混把图片读取放在__init__里结果数据集一大内存直接爆掉或者启动时间长得离谱。我的做法是在__init__里完成这几件事第一扫描目录构建rainy_paths和clean_paths两个列表并且保证它们一一对应。如果是“一对多”结构rainy_paths会比clean_paths长但每个rainy_path都能通过映射找到对应的clean_path。第二解析场景 ID建立scene_to_indices字典方便后续按场景划分。第三接收transform参数。注意去雨任务的 transform 有个特殊要求输入和标签必须做完全相同的几何变换。如果你用torchvision.transforms里的RandomCrop分别作用在两张图上很可能裁到不同位置配对就废了。正确做法是用torchvision.transforms.functional或者自己写一个同时接收两张图的 transform 函数。第四可选地接收一个cache参数。如果数据集不大比如干净图只有几百张可以在__init__里把干净图预加载到内存带雨图按需读取。这样能显著减少 IO因为干净图会被反复读取。这里有个经验不要用os.listdir直接拿文件列表因为它不保证顺序。如果你先listdir带雨图目录再listdir干净图目录两个列表的顺序可能不一致导致配对错乱。稳妥的做法是用sorted(os.listdir(...))或者自己写一个按场景 ID 排序的逻辑。2.2__getitem__的配对逻辑与异常处理__getitem__是核心它接收一个索引返回一对张量。我的实现大致是这样的思路def __getitem__(self, idx): rainy_path self.rainy_paths[idx] scene_id self.parse_scene_id(rainy_path) clean_path self.scene_to_clean[scene_id] rainy_img Image.open(rainy_path).convert(RGB) clean_img Image.open(clean_path).convert(RGB) # 同步几何变换 rainy_tensor, clean_tensor self.joint_transform(rainy_img, clean_img) return rainy_tensor, clean_tensor看起来简单但有几个坑必须处理。坑一图片尺寸不一致。有些数据集的带雨图和干净图分辨率不同比如带雨图是裁剪过的干净图是原图。直接送进网络会报错。我的做法是在__init__里做一次尺寸检查如果发现不一致要么统一 resize要么在__getitem__里做中心裁剪。中心裁剪更安全因为它不会引入插值伪影而去雨任务对伪影很敏感。坑二图片模式不是 RGB。有些干净图是灰度图或者带 alpha 通道的 PNG。convert(RGB)能解决大部分问题但如果原图是 16 位深度convert会截断成 8 位损失信息。如果你的数据集是 16 位需要用numpy读取并手动归一化。坑三文件损坏或缺失。大规模数据集里偶尔会有损坏的图片。如果__getitem__直接抛异常训练会中断。我的做法是用try-except包住读取逻辑出错时随机返回另一对样本并打印警告。这样训练不会断但你要定期检查警告日志看看是不是有系统性问题。坑四配对失败。如果某个带雨图的场景 ID 在干净图里找不到说明数据集本身有问题。这种情况我建议在__init__阶段就检查并报错而不是等到训练时才发现。可以在__init__里遍历所有带雨图验证每个场景 ID 都有对应的干净图缺失的直接列出来。2.3 同步几何变换去雨任务最容易写错的地方这一小节单独拿出来讲因为它太重要了。去雨任务的输入和标签是同一场景的两张图空间位置必须严格对齐。如果你对输入做随机裁剪对标签做另一套随机裁剪两张图就错位了模型学到的是错误映射loss 可能下降但效果完全不对。正确的做法是生成一组随机参数然后把这组参数同时应用到两张图上。比如随机裁剪你先随机生成裁剪的左上角坐标和宽高然后用同一组坐标去裁两张图。随机翻转同理生成一个布尔值两张图一起翻或不翻。torchvision.transforms.functional提供了crop、hflip、vflip、rotate等函数它们接收具体的参数而不是随机生成参数正好适合这种场景。你可以自己封装一个JointTransform类在__call__里生成随机参数并应用到两张图。如果你用的是albumentations库它原生支持additional_targets可以把标签图作为额外目标保证变换同步。配置方式是在Compose里加additional_targets{image0: image}然后调用时传入imagerainy, image0clean。这个库在处理同步变换上比手写更省心但要注意版本兼容性不同版本的 API 有差异。注意不要用RandomResizedCrop这类会改变图像内容的变换去做同步因为两张图如果内容不同一张有雨一张没雨resize 后的插值结果会有细微差异虽然影响不大但严格来说不如直接裁剪干净。我通常只用随机裁剪 随机翻转最多加一个 90 度旋转。3. 数据加载性能为什么你的 GPU 利用率上不去模型训练慢很多人第一反应是模型太大或者 GPU 不行。但在去雨任务里更常见的原因是数据加载成了瓶颈。去雨模型的输入输出都是高分辨率图片读取、解码、变换的开销很大如果 DataLoader 配置不当GPU 大部分时间在等数据。我做过一个对比实验同样的模型、同样的数据量只调整 DataLoader 的参数训练速度差了将近一倍。下面把我踩过的坑和调优经验整理出来。3.1num_workers不是越大越好num_workers控制数据加载的并行进程数。理论上越大越快但实际上受限于 CPU 核心数、磁盘 IO 和内存带宽。我试过在 8 核机器上设num_workers16结果反而比num_workers4慢因为进程切换和内存拷贝的开销超过了并行收益。我的经验值是num_workers设为 CPU 物理核心数的 1 到 1.5 倍。比如 8 核就设 8 或 12。如果是 SSD可以适当调大如果是机械硬盘调大反而会因为随机读取导致寻道开销增加。另外num_workers 0时Dataset 里的对象会被复制到子进程如果你在__init__里缓存了大量数据每个 worker 都会复制一份内存占用会成倍增加。解决办法是用torch.utils.data.DataLoader的persistent_workersTrue配合prefetch_factor或者把缓存改成共享内存。3.2pin_memory和prefetch_factor的配合pin_memoryTrue会把数据放到锁页内存加速 CPU 到 GPU 的传输。这个参数在 GPU 训练时几乎必开但要注意它和num_workers的配合。如果num_workers0pin_memory的效果有限因为数据加载和传输是串行的。prefetch_factor控制每个 worker 预取多少个 batch。默认是 2对于去雨这种 IO 密集的任务可以调到 4 或 6。但调太大也会占内存需要根据你的图片分辨率和 batch size 权衡。我通常的配置是loader DataLoader( dataset, batch_size8, shuffleTrue, num_workers8, pin_memoryTrue, prefetch_factor4, persistent_workersTrue, drop_lastTrue )drop_lastTrue在去雨任务里建议开启因为最后一个 batch 如果只有一张图BatchNorm 的统计会不稳定影响训练。3.3 图片解码的优化从 PIL 到 OpenCV 再到 DALIPIL 是 PyTorch 生态里最常用的图片读取库但它的解码速度一般。如果你的数据集很大可以换成 OpenCV 的cv2.imread速度通常快 2 到 3 倍。注意 OpenCV 读进来是 BGR 顺序需要转成 RGB。更进一步如果 CPU 解码仍然是瓶颈可以考虑 NVIDIA 的 DALI 库它能把解码放到 GPU 上做。但 DALI 的配置比较复杂而且对自定义变换的支持不如 PIL 灵活。我的建议是先用 PIL 跑通确认瓶颈在数据加载后再考虑换库。很多时候把num_workers和prefetch_factor调好问题就解决了。还有一个容易被忽略的点图片格式。PNG 是无损压缩解码慢但质量高JPEG 解码快但有压缩伪影。去雨任务对伪影敏感我建议训练时用 PNG如果磁盘空间不够可以用高质量的 JPEGquality 95 以上。我实测过quality 95 的 JPEG 和 PNG 在去雨模型上的效果差异很小但读取速度快不少。4. 训练中的那些坑从 loss 不下降到雨痕残留数据和加载都搞定之后训练过程本身也有一堆坑。去雨任务和普通图像恢复任务相比有几个特殊之处雨痕是叠加在图像上的不是简单遮挡雨痕有方向性不同方向的雨需要不同的处理雨痕的密度变化很大从稀疏雨滴到暴雨都有。这些特性决定了你的训练策略不能照搬超分或者去噪。4.1 损失函数的选择L1、L2 还是感知损失L2 损失MSE在去雨任务里有个著名的问题它倾向于产生模糊的结果。因为 L2 对误差的惩罚是平方的模型会倾向于预测所有可能值的平均导致雨痕区域被“抹平”而不是“去除”。你去看用纯 L2 训练的去雨模型雨痕是没了但原本的纹理也糊了。L1 损失比 L2 好一些因为它对异常值的惩罚是线性的不会过度平滑。但 L1 也有问题它在零点不可导训练后期容易震荡。实践中我通常用SmoothL1Loss也叫 Huber Loss它在误差小的时候表现为 L2误差大的时候表现为 L1兼顾了稳定性和锐度。如果你追求更好的视觉效果可以加感知损失Perceptual Loss用预训练的 VGG 提取特征比较去雨结果和干净图在特征空间的差异。感知损失能让结果更符合人眼观感但训练成本会上升而且需要小心权重设置感知损失权重太大时PSNR 会下降。我的经验配置是L1 0.1 * Perceptual感知损失用 VGG16 的 relu2_2 和 relu4_3 层。这个配置在多个数据集上表现稳定。4.2 雨痕残留和过度去雨的平衡训练去雨模型时你会遇到两种相反的问题一种是雨痕没去干净另一种是把原本的纹理当成雨去掉了导致图像过度平滑。雨痕残留通常发生在训练早期模型还没学会识别雨痕。这时候可以检查一下学习率是不是太大或者数据增强是不是太激进。如果雨痕是细密的斜线而你的随机裁剪把大部分雨痕裁掉了模型就学不到。解决办法是确保裁剪后的区域仍然包含足够的雨痕可以通过统计雨痕像素比例来筛选。过度去雨通常发生在训练后期模型为了降低 loss把一些高频纹理也当成雨去掉了。这时候可以降低学习率或者引入一个“保真度”约束比如在干净图的高频区域加大 loss 权重。另一个技巧是使用对抗损失GAN Loss让判别器判断结果是否像真实干净图但这会显著增加训练复杂度新手不建议一上来就用。4.3 验证集指标和肉眼观感的差距PSNR 和 SSIM 是去雨任务最常用的指标但它们和肉眼观感经常不一致。我遇到过 PSNR 很高的模型结果图看起来还是有雾感也遇到过 PSNR 一般但观感很干净的模型。原因是 PSNR 对亮度偏移很敏感如果模型整体偏亮或偏暗PSNR 会掉但人眼可能觉得还好。SSIM 好一些但它对纹理的敏感度也不够。我的做法是训练时看 PSNR/SSIM但最终选模型一定要肉眼对比。我会在验证集里挑几张有代表性的图稀疏雨、密集雨、不同背景把不同 epoch 的结果拼在一起看。有时候 PSNR 最高的 epoch 并不是观感最好的选那个观感好的指标差 0.1 到 0.2 完全可以接受。另外验证时要注意把结果保存成图片不要只看数值。我习惯每几个 epoch 就导出一次验证结果用图片浏览器快速翻看比盯着数字直观得多。5. 从单张推理到批量部署实际使用中的工程问题模型训练好之后真正的挑战才开始。研究阶段你处理的是固定尺寸的图片实际使用中你会遇到各种尺寸、各种来源的图片。这一节讲讲工程化落地时的问题。5.1 任意尺寸输入的推理策略去雨模型通常是全卷积网络理论上支持任意尺寸输入。但实际推理时如果图片太大显存会爆如果图片太小感受野不够效果会下降。我的策略是把大图切成有重叠的 patch分别推理后拼接。重叠区域取平均避免拼接缝。patch 大小根据你的模型和显存来定我通常用 512x512重叠 64 像素。这个配置在 8G 显存的卡上跑得很稳。切 patch 时要注意不要简单地从左上角开始切因为雨痕可能正好落在切割线上。有重叠的切法能缓解这个问题但也不能完全避免。如果对质量要求极高可以用滑动窗口步长小于 patch 大小重叠区域用高斯权重融合。5.2 批量推理时的显存管理批量推理时显存占用和 batch size 成正比。如果你要处理几千张图不要一次性全部加载而是用 DataLoader 流式读取推理完一张释放一张。另外推理时记得用torch.no_grad()否则会保存计算图显存占用会大很多。如果模型里有 BatchNorm推理时要设model.eval()否则会使用当前 batch 的统计量结果不稳定。我见过有人推理时忘了eval()结果单张推理和批量推理的结果不一致排查了半天。这个坑很隐蔽因为训练时一切正常只有推理时才暴露。5.3 结果后处理锐化和色彩校正模型输出的图片有时候会偏软或者偏色。偏软是因为 L1/L2 损失的平滑效应偏色是因为训练数据的色彩分布和实际图片不一致。后处理可以缓解这些问题。锐化用简单的UnsharpMask就行强度不要太大否则会放大噪声。色彩校正可以用直方图匹配把结果的色彩分布对齐到输入图注意是对齐到输入图不是干净图因为实际使用时你没有干净图。但后处理要谨慎它可能掩盖模型本身的问题。我的建议是先看原始输出如果问题不严重不要加后处理如果确实需要把后处理参数记录下来方便复现。6. 一些零散但重要的经验最后分享几个我在做去雨项目时积累的零散经验它们不属于某个特定环节但每个都帮我省过时间。关于数据增强去雨任务的增强要克制。随机裁剪和翻转是安全的但颜色抖动要小心因为颜色变化可能让模型把颜色异常误认为雨痕。我通常只做轻微的亮度调整幅度控制在 0.1 以内。关于预训练模型如果你的去雨网络用了预训练 backbone比如 ResNet加载权重时要注意有些预训练模型是在 ImageNet 上训练的输入归一化参数不同。用错归一化参数会导致训练初期 loss 爆炸。关于随机种子去雨任务的随机性来源很多数据划分、数据增强、权重初始化固定种子能让你复现实验。但要注意num_workers 0时DataLoader 的随机种子需要在 worker 初始化函数里设置否则每个 worker 的随机序列可能相同导致增强结果重复。关于日志训练去雨模型时建议把每个 epoch 的验证结果图保存下来按 epoch 编号命名。这样训练结束后可以快速回看模型是怎么逐步变好的也能发现某些 epoch 的异常比如突然变差又恢复可能是学习率调度的问题。关于显存去雨模型的显存占用和图片分辨率是平方关系。512x512 的图batch size 8 可能占 6G 显存1024x1024 的图同样的 batch size 可能就要 20G 以上。如果你的显存不够优先降 batch size而不是降分辨率因为分辨率对去雨效果的影响更大。关于评估不要只在合成数据集上评估。合成雨和真实雨在物理特性上有差距合成雨通常是均匀的斜线真实雨有雨滴、有雾、有镜头上的水渍。如果你的模型只在合成数据上训练放到真实照片上效果会打折扣。有条件的话收集一些真实雨天照片做定性评估哪怕没有干净图对照肉眼看看雨痕去得干不干净也是有价值的。我在实际项目里最深的一个体会是去雨任务的上限往往不在模型结构而在数据质量。一个简单的 U-Net配上高质量、多样化的成对数据效果可能比一个复杂网络配劣质数据好得多。所以如果你刚开始做去雨先把数据集的配对、划分、加载这三件事做扎实比急着换模型架构重要得多。
