1. 图像去雨到底在解决什么问题
先聊点实际的。图像去雨,就是给定一张带雨纹的图,让模型学会把这层“干扰”剥掉,恢复出干净背景。这个任务听起来简单,但做起来比想象中麻烦得多——雨不是均匀撒在画面上的,它有方向、有粗细、有透明度,近处雨滴清晰得像短线,远处雨丝模糊成薄雾,再加上背景纹理和雨纹混在一起,算法很难区分哪些细节该保留、哪些该抹掉。
我写这篇博文的目的很直接:不绕弯子,不堆理论,手把手把一套完整的 pytorch 图像去雨实现代码写出来,从数据处理、网络设计到训练推理,每行注释都写上为什么这么写。适合两类人看:一是刚入门深度学习、想找个小任务练手跑通全流程的同学;二是已经跑过分类或检测、想转图像恢复方向但没有现成代码参考的开发者。看完你应该能拿着这套框架,自己改数据、换网络,去跑自己的去雨实验。
为什么会选去雨这个任务?因为它对新手非常友好。输入是图像,输出是图像,不需要做目标框、不需要做分割mask这种复杂的标注任务,数据配对方式直接,网络结构也不难理解。但反过来,它又比单纯分类要难出不少——分类只需要判别一个标签,去雨要逐像素重建图像,这涉及特征提取、残差学习、损失设计等多个环节,做完这一个项目,很多基础能力都练到了。
2. 环境准备与依赖安装
2.1 pytorch 环境搭建要点
先说环境。很多人卡在第一步,pytorch 装了半天跑不起来,其实大部分问题都出在版本不匹配上,尤其是 GPU 版和 CUDA 的对应关系。
我的建议是项目开始时不要用 conda 默认源,直接用官方命令生成器去选版本。打开 pytorch 官网,根据自己的操作系统、包管理工具和 CUDA 版本,它会直接给出安装命令。如果你没有特别需求,CUDA 就选稳定版的,不要追最新,太新的 CUDA 版本有时候配套的驱动要求也高,反而容易出问题。
如果你用的是 CPU 机器,那更省事,直接装 CPU 版就行。很多人纠结“CPU 能不能跑图像任务”,说实话能跑,但确实慢。一个 200 轮训练的小模型,GPU 半小时跑完的话,CPU 可能要一晚上。建议至少用个带 CUDA 的显卡,哪怕显存只 4G,也足够跑我们的轻量模型。
我实际用的安装命令大概是这样的,大家可以参考:
conda create -n derain python=3.9 conda activate derain pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118 pip install numpy opencv-python pillow matplotlib tqdm这里提一个关键点:torchvision和torch的版本必须匹配。很多人单独 pip 装 torch 成功了,后来装 torchvision 又强制升了某个依赖,结果两个库版本对不上,直接 import 报错。推荐一条命令同时装,让 pip 自己解析依赖关系。
2.2 常见安装坑
我在多个环境里配过 pytorch,遇到最多的一个报错是OSError: [WinError 1114] 动态链接库(DLL)初始化例程失败,这个报错后面通常还跟着error loading "torch/lib/c10.dll" or one of its dependencies。什么意思呢?就是 torch 加载 C++ 扩展库时,一个依赖的 DLL 找不到或者版本不对。最常见的原因是机器上缺少 Microsoft Visual C++ Redistributable,或者安装的是老版本。去微软官网下载最新的vc_redist.x64.exe装上,基本能解决。
还有一类常见问题是 CUDA 装好了但 torch 检测不到 GPU。这时候别急着重装,先检查两件事:第一,打开命令行执行nvidia-smi,看驱动版本是否支持你的 CUDA 版本;第二,用python -c "import torch; print(torch.cuda.is_available())"验证 torch 的可用状态。如果返回 False,大概率是装错了 CUDA 版本的包,卸载重装正确版本就行。
3. 数据准备与预处理
3.1 去雨数据集怎么组织
去雨训练需要成对的图像:一张是带雨的输入图,一张是同一个场景的干净图。公开数据集有不少,比如 Rain100H、Rain100L、DidMDN 等,因为版权问题我不直接给下载链接,大家去网上搜索数据集名称都能找到。
但我想多说一句数据目录的组织方式,这直接影响后面 DataLoader 的写法。比较推荐的结构是这样:
data/ train/ input/ 001.png 002.png target/ 001.png 002.png val/ input/ 001.png target/ 001.png就是说,把所有带雨图放在input文件夹,干净图放在target文件夹,两个文件夹里的文件名一一对应。这样加载数据的时候,只需要拿到文件名列表,按名字去两个目录分别读图就行。不用搞什么 CSV、JSON 映射文件,简单直接,还不容易出错。
3.2 数据加载与增强
写一个自定义 Dataset 类,继承torch.utils.data.Dataset。核心方法是__getitem__,它返回一对数据:带雨的输入图和对应的干净图。这里面有几个细节值得注意。
第一,图像读取问题。OpenCV 读出来的是 BGR 通道顺序,而 pytorch 训练时用的基本是 RGB。如果不做转换,模型会学得很奇怪——颜色失真但 loss 还很低,因为模型学会了在 BGR 空间里做映射。所以每次读图后要手动cv2.cvtColor(img, cv2.COLOR_BGR2RGB)。
第二,归一化问题。图像像素值范围是 0 到 255,直接喂给网络会让梯度很大,训练不稳定。通常要除以 255,缩放到 0 到 1。另一个方案是计算整个数据集的均值和标准差做标准化,但对去雨这种逐像素重建任务,直接用 0-1 范围就够了,很多论文也是这么干的。
第三,数据增强。训练时做随机裁剪、随机翻转、随机旋转,可以显著提高模型泛化能力。我常用的增广策略是:随机裁剪成 256x256 的块,50% 概率水平翻转,50% 概率垂直翻转。实现的时候注意一点:输入图和目标图要做完全相同的变换,不能一个翻转了另一个没翻转,否则训练就废了。所以随机数种子要统一或者用固定的随机变量控制。
下面是一个参考实现:
import torch from torch.utils.data import Dataset import cv2 import numpy as np import os class DerainDataset(Dataset): def __init__(self, input_dir, target_dir, crop_size=256, train=True): self.input_dir = input_dir self.target_dir = target_dir self.crop_size = crop_size self.train = train # 只需要列出一个目录的文件名,另一个目录按名字取对应图 self.names = sorted(os.listdir(input_dir)) def __len__(self): return len(self.names) def __getitem__(self, idx): name = self.names[idx] # 读取输入图和目标图,注意顺序转换 BGR -> RGB img_in = cv2.imread(os.path.join(self.input_dir, name)) img_tar = cv2.imread(os.path.join(self.target_dir, name)) img_in = cv2.cvtColor(img_in, cv2.COLOR_BGR2RGB) img_tar = cv2.cvtColor(img_tar, cv2.COLOR_BGR2RGB) # 统一尺寸:不想裁剪的话就 resize,但训练时裁剪更好 # 如果是验证集,直接缩放到固定大小 if not self.train: img_in = cv2.resize(img_in, (256, 256), interpolation=cv2.INTER_CUBIC) img_tar = cv2.resize(img_tar, (256, 256), interpolation=cv2.INTER_CUBIC) else: # 随机裁剪出 crop_size 大小的块 h, w = img_in.shape[:2] if h < self.crop_size or w < self.crop_size: img_in = cv2.resize(img_in, (self.crop_size + 1, self.crop_size + 1), interpolation=cv2.INTER_CUBIC) img_tar = cv2.resize(img_tar, (self.crop_size + 1, self.crop_size + 1), interpolation=cv2.INTER_CUBIC) h, w = img_in.shape[:2] y = np.random.randint(0, h - self.crop_size) x = np.random.randint(0, w - self.crop_size) img_in = img_in[y:y+self.crop_size, x:x+self.crop_size] img_tar = img_tar[y:y+self.crop_size, x:x+self.crop_size] # 随机翻转:用同一个随机数控制输入和目标的翻转方向 if self.train: if np.random.rand() > 0.5: img_in = cv2.flip(img_in, 1) img_tar = cv2.flip(img_tar, 1) if np.random.rand() > 0.5: img_in = cv2.flip(img_in, 0) img_tar = cv2.flip(img_tar, 0) # 归一化到 [0, 1],并转为 tensor,维度从 H,W,C 变为 C,H,W img_in = torch.from_numpy(img_in.astype(np.float32) / 255.0).permute(2, 0, 1) img_tar = torch.from_numpy(img_tar.astype(np.float32) / 255.0).permute(2, 0, 1) return img_in, img_tar有人可能会问,为什么验证集用resize而不是也做随机裁剪?原因很简单:验证集要衡量模型在完整图像上的效果,如果做随机裁剪,每次验证只看到局部信息,评估结果不稳定。所以验证集固定缩放到统一尺寸,保证可复现。
3.3 DataLoader 的配置细节
有了 Dataset 之后,DataLoader 的配置也有讲究。核心参数:
train_loader = torch.utils.data.DataLoader( train_dataset, batch_size=8, shuffle=True, num_workers=2, pin_memory=True ) val_loader = torch.utils.data.DataLoader( val_dataset, batch_size=1, shuffle=False, num_workers=1, pin_memory=True )shuffle=True是必须的,训练集不打乱顺序的话,模型会学到数据的排列规律,导致验证集表现下降。pin_memory=True在 GPU 机器上能加快数据从 CPU 到 GPU 的拷贝速度,值得加上。
4. 模型设计:从零手写一个可用的去雨网络
4.1 网络结构设计思路
很多刚接触去雨的人会想:直接拿 ResNet 改一改不就行了?能行,但效果不会特别好。原因在于,去雨任务输出是逐像素的图像,需要网络在多个尺度上都能抓住纹理信息——小雨纹要精细地去除,大雨滴又带有模糊和透明度变化,单靠一个固定感受野的深层网络很难两头兼顾。
我采用的思路是保留主干的残差设计,加上一个重要的处理手段:在第一层就预测“雨纹残差”,而不是直接预测干净图。这是什么意思呢?我们不直接让网络输出y = clean_image,而是让网络输出y = rainy_image - clean_image,也就是雨纹本身。最后恢复的时候,用带雨图减去预测的雨纹,得到去雨结果。
这个设计的巧妙之处在于,残差学习降低了网络的学习难度。干净图像和带雨图像有大量共享信息(背景纹理、边缘、色彩),如果直接让网络输出干净图,它需要“记住”全部背景细节,这不仅浪费参数,还容易把细节改坏。但如果只输出“差异”,网络只需要专注于捕捉雨纹特征,背景信息不用重建,输入图里本来就有,做一次减法就拿回来了。
还有一个细节:雨纹通常是非线性的,亮度高、边缘尖锐,所以网络的最后一个激活函数不能是 ReLU 或者 Sigmoid,因为它们会对输出做非线性压缩或截断,导致雨纹估计不准确。这里直接用线性输出就行,让网络自己学输出范围。
4.2 逐段代码与注释
下面是我们模型的完整实现,我尽量把每一层的意图都写清楚:
import torch import torch.nn as nn class BasicBlock(nn.Module): """ 基础残差块。结构:Conv -> BN -> ReLU -> Conv -> BN,然后与输入相加。 相比直接堆叠卷积层,残差连接让梯度能顺畅地传到浅层, 训练时不容易出现梯度消失,网络也能加深到更多层。 """ def __init__(self, channels): super(BasicBlock, self).__init__() self.conv1 = nn.Conv2d(channels, channels, kernel_size=3, padding=1, bias=False) self.bn1 = nn.BatchNorm2d(channels) self.relu = nn.ReLU(inplace=True) self.conv2 = nn.Conv2d(channels, channels, kernel_size=3, padding=1, bias=False) self.bn2 = nn.BatchNorm2d(channels) def forward(self, x): identity = x out = self.conv1(x) out = self.bn1(out) out = self.relu(out) out = self.conv2(out) out = self.bn2(out) out += identity # 残差连接 out = self.relu(out) return out class DerainNet(nn.Module): """ 去雨网络主体: 1. 先用一个卷积层把 3 通道输入映射到高维特征空间 2. 经过多个残差块提取雨纹特征 3. 通过一个卷积层输出雨纹残差 4. 用输入图减去雨纹残差,得到去雨图 """ def __init__(self, in_channels=3, hidden_channels=32, num_blocks=8): super(DerainNet, self).__init__() # 第一层:将 RGB 三通道映射到特征空间。 # 为什么用 3x3 卷积而不是更大的?3x3 足够捕获局部纹理信息, # 而且参数少,堆叠多个 3x3 卷积可以获得与大卷积核相等的感受野。 self.entry = nn.Sequential( nn.Conv2d(in_channels, hidden_channels, kernel_size=3, padding=1, bias=False), nn.ReLU(inplace=True) ) # 堆叠残差块。残差块数量越多,网络的感受野越大,能捕获更大范围的雨纹特征。 # 但也不是越多越好,太多层会导致训练变慢,而且小数据集上容易过拟合。 blocks = [] for _ in range(num_blocks): blocks.append(BasicBlock(hidden_channels)) self.body = nn.Sequential(*blocks) # 输出层:从特征空间映射回 3 通道残差。 # 注意这里没有激活函数,因为残差值可正可负, # 加个 Sigmoid 或 ReLU 反而会限制输出范围。 self.output = nn.Conv2d(hidden_channels, in_channels, kernel_size=3, padding=1, bias=False) def forward(self, x): feature = self.entry(x) feature = self.body(feature) rain_residual = self.output(feature) # 关键操作:输入减去预测的雨纹残差,得到干净图像 output = x - rain_residual return output这段代码加起来不过 70 行,但它已经是一个完整可训练的去雨网络了。第一次写的时候我犯过一个错误:在entry层后面没有加ReLU,导致特征值可能为负,而且后面的残差块作用在负值上,训练时 loss 降得很慢。后来加上了ReLU做非线性映射,问题就消失了。
4.3 为什么这样设计
选择 32 个隐藏通道和 8 个残差块,是一个平衡性能和速度的决定。通道数太少,特征表达不够,雨纹去不干净;通道数太多,参数爆炸,训练时间长。我试过 64 通道版本,效果提升有限,但显存占用和训练速度明显变差。对于 256x256 的输入,32 通道版本在 8G 显存上可以跑 batch_size=16,而 64 通道最多跑 batch_size=8。
至于残差块数量,8 个块在感受野上已经覆盖了足够大的范围。一个 3x3 卷积的感受野是 3x3,两个叠加是 5x5,三个叠加是 7x7,8 个残差块加上前面的入口卷积,感受野已经超过 20x20,对常见的雨滴尺寸来说完全够用。如果数据集里有特别粗的大雨滴,可以适当增加到 12 个块,但不要盲目加到 20 个以上,收益递减明显。
5. 训练流程:损失函数、优化器与完整训练循环
5.1 损失函数选择
去雨任务常用的损失函数有三种:L1 损失、L2 损失、感知损失。
L2 损失也就是 MSE,最开始入门时我用的是它,但后来发现一个问题:L2 对离群点非常敏感,像素误差较大的地方会主导梯度,导致训练不稳定。而且 L2 倾向于生成模糊的结果,因为这个损失会平均掉多种可能的锐利细节。
L1 损失,也就是平均绝对误差,它不会放大离群点的梯度,训练更稳定,生成的图像也相对锐利。现在很多恢复任务的首选都是 L1。
感知损失需要额外的预训练网络提取特征,计算复杂度高,而且对新手不太友好,暂时不推荐。
我最终的方案是L1 损失 + 0.05 * MSE 损失,用 L1 主导训练、MSE 做微调。这个组合在实际实验中效果比单用任何一项都要好,因为 L1 保证全局结构稳定,MSE 进一步约束像素级别的精确度。
5.2 训练循环代码
训练循环看起来很标准化,但里面有几个容易被忽视的细节,我都写在注释里:
import torch import torch.nn as nn import torch.optim as optim from tqdm import tqdm def train_one_epoch(model, dataloader, optimizer, criterion, device): model.train() # 切换到训练模式,这会启用 BN 的滑动平均更新和 Dropout total_loss = 0.0 for inputs, targets in tqdm(dataloader, desc="Training"): inputs = inputs.to(device) targets = targets.to(device) optimizer.zero_grad() outputs = model(inputs) loss = criterion(outputs, targets) # loss 反向传播之前,检查是否为 NaN。出现 NaN 通常意味着学习率过大 # 或输入数据里存在异常值,尽早发现比训练完了再排查要高效得多。 if torch.isnan(loss): print("Loss is NaN, stopping epoch") break loss.backward() # 梯度裁剪:防止梯度爆炸。尤其在小数据集上训练较深网络时, # 偶尔会出现单批数据产生特别大的梯度,不加裁剪会把参数冲坏。 torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=0.5) optimizer.step() total_loss += loss.item() return total_loss / len(dataloader) def evaluate(model, dataloader, criterion, device): model.eval() # 切换到评估模式,BN 使用累积的统计量,Dropout 关闭 total_loss = 0.0 with torch.no_grad(): # 评估阶段不计算梯度,大幅减少内存占用 for inputs, targets in dataloader: inputs = inputs.to(device) targets = targets.to(device) outputs = model(inputs) loss = criterion(outputs, targets) total_loss += loss.item() return total_loss / len(dataloader)主训练流程:
device = torch.device("cuda" if torch.cuda.is_available() else "cpu") model = DerainNet().to(device) # 优化器选择。Adam 是默认选择,它自带自适应学习率, # 对新手非常友好,基本不需要手动调整太多参数。 # 学习率初始值设为 1e-3,这是一个比较稳妥的起点。 optimizer = optim.Adam(model.parameters(), lr=1e-3) # 损失函数:L1 主导 + 0.05 * MSE 辅助 criterion = lambda pred, gt: nn.functional.l1_loss(pred, gt) + 0.05 * nn.functional.mse_loss(pred, gt) # 学习率调度:每 20 轮衰减为原来的 0.7 倍 scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=20, gamma=0.7) EPOCHS = 100 best_loss = float('inf') for epoch in range(EPOCHS): train_loss = train_one_epoch(model, train_loader, optimizer, criterion, device) val_loss = evaluate(model, val_loader, criterion, device) scheduler.step() print(f"Epoch {epoch+1}/{EPOCHS} - Train Loss: {train_loss:.6f}, Val Loss: {val_loss:.6f}") # 保存最优模型,用验证 loss 作为判断标准而不是训练 loss, # 因为训练 loss 持续下降往往伴随着过拟合,验证 loss 更能反映真实泛化能力。 if val_loss < best_loss: best_loss = val_loss torch.save(model.state_dict(), "best_derain_model.pth")5.3 训练技巧
我踩过的几个实质性的坑,在这里集中说一下。
第一个坑是关于model.train()和model.eval()的切换。前面代码里已经正确写了,但很多人会在评估的时候忘记加model.eval(),导致 BN 层还在用当前 batch 的统计量,验证集 loss 一会高一会低,非常不稳定。而训练时忘记切换回来,又会导致 BN 的滑动平均更新不连贯。这个错误非常隐蔽,因为模型仍然能“跑通”,就是结果很差。
第二个坑是学习率调度。刚开始训练时学习率设到 0.001,跑几十个 epoch 后如果 loss 不再下降,就手动或通过 scheduler 把学习率降下去,一般还能再降一个台阶。如果从头到尾不调整学习率,后期 loss 会在一个平台值附近震荡,很难有实质提升。
第三个坑是 batch size 对 BN 的影响。如果用 BatchNorm,batch size 不能太小,否则 BN 统计量估计不准。我试过 batch_size=2 跑这个网络,效果明显变差。如果你显存不够,要么降低图像裁剪尺寸,要么改用 GroupNorm 替代 BatchNorm,效果会稳定很多。
6. 推理与评估:如何量化去雨效果
6.1 评估指标 PSNR 和 SSIM
训练模型不是终点,关键是看模型在新图上的表现。但“看起来干净”太主观,我们需要两个量化指标:PSNR 和 SSIM。
PSNR,全称峰值信噪比,衡量像素级别的恢复精度。计算公式是PSNR = 10 * log10(MAX^2 / MSE),其中MAX是像素最大值(图像范围 0-1 时是 1)。这个值越大越好,通常在 25 到 40 之间,高于 35 肉眼看起来就比较接近原图了。
SSIM,全称结构相似性指数,衡量图像纹理结构的一致性,取值 0 到 1,越接近 1 越好。PSNR 只盯着像素差异,SSIM 还考虑了亮度、对比度和结构三方面信息,两者结合才能比较全面地评价去雨效果。
6.2 推理代码
推理代码比训练简单多了,核心就是把模型加载出来,对输入图前向推理一次:
import torch import cv2 import numpy as np def inference(model_path, input_image_path, output_image_path, device): model = DerainNet().to(device) # 加载保存的权重。注意 load_state_dict 要设置 map_location, # 否则在 GPU 上训练、CPU 上推理时会因为键值不匹配或设备不兼容而报错。 model.load_state_dict(torch.load(model_path, map_location=device)) model.eval() # 读取并预处理图像 img = cv2.imread(input_image_path) img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) # 转换为 tensor,加上 batch 维度 img_tensor = torch.from_numpy(img_rgb.astype(np.float32) / 255.0).permute(2, 0, 1).unsqueeze(0) img_tensor = img_tensor.to(device) with torch.no_grad(): output = model(img_tensor) # 从 tensor 转回图像数组,注意去掉 batch 维度、调整通道顺序 output = output.squeeze(0).permute(1, 2, 0).cpu().numpy() output = np.clip(output, 0.0, 1.0) # 确保像素值在合法范围内 output = (output * 255.0).astype(np.uint8) output_bgr = cv2.cvtColor(output, cv2.COLOR_RGB2BGR) cv2.imwrite(output_image_path, output_bgr) print(f"Saved result to {output_image_path}")计算 PSNR 和 SSIM 的代码:
from skimage.metrics import peak_signal_noise_ratio, structural_similarity def compute_metrics(pred_path, gt_path): pred = cv2.imread(pred_path) gt = cv2.imread(gt_path) # 注意两张图的尺寸必须一致,否则指标计算会报错 if pred.shape != gt.shape: pred = cv2.resize(pred, (gt.shape[1], gt.shape[0])) # PSNR 计算可以直接用在 BGR 图像上,但 SSIM 需要指定 win_size psnr = peak_signal_noise_ratio(gt, pred) ssim = structural_similarity(gt, pred, multichannel=True, win_size=7) return psnr, ssim如果不想依赖 skimage,PSNR 也可以自己用 tensor 计算,几行代码就搞定:
def compute_psnr(pred, gt, max_val=1.0): mse = torch.mean((pred - gt) ** 2) psnr = 10 * torch.log10(max_val * max_val / mse) return psnr.item()6.3 一个需要注意的形状问题
推理时经常遇到一个问题:训练时用固定尺寸的裁剪块,但测试图像尺寸是任意的。如果网络全卷积化,理论上任意尺寸都能跑,但某些操作如 BatchNorm 在推理时对任意尺寸依然有效。不过如果你训练和验证都用固定 256,推理大图时可能效果没那么好。一个简单的处理是先缩放到 256 的倍数,推理完再缩放回去,能减少拼接伪影。
7. 常见问题排查与优化建议
7.1 常见错误一览
实际跑这个项目时,下面这些问题我基本都遇到过,整理成一个速查表,方便排查:
| 现象 | 可能原因 | 解决方案 |
|---|---|---|
| Loss 一直是 NaN | 学习率过大、输入数据存在 NaN 像素 | 调低学习率,检查数据归一化是否正确 |
| 训练 loss 下降,验证 loss 不断升高 | 过拟合 | 增加数据增强、减少 epoch、增加训练数据 |
| 输出的去雨图有棋盘格伪影 | 反卷积使用不当 | 用亚像素卷积或直接上采样 + 卷积替代 |
| 图像整体变暗或偏色 | 归一化方式不对,模型输出范围偏移 | 检查是否减掉了均值而没有加回来 |
| 测试时显存不足 | batch size 过大或输入图像过大 | 减小 batch size,或对输入做切片推理 |
| 同一张图多次推理结果不同 | 模型处于训练模式,没有调用 model.eval() | 推理前必须加 model.eval() |
| 加载权重时报尺寸不匹配 | 模型结构和保存权重时不一致 | 检查隐藏层通道数、残差块数量是否一致 |
7.2 如何进一步提升模型效果
当前实现的模型是一个基础版本,能在小数据集上跑通并给出不错的结果。如果想让效果更上一层楼,有几个方向可以尝试。
方向一:金字塔结构。大多数雨滴大小不一,单一尺度的特征提取很难同时适应。可以在网络中间加入下采样和上采样分支,形成编码器-解码器结构,让网络在不同尺度上学习雨纹特征。比如输入先下采样到 1/2 和 1/4,分别提取特征后再逐级融合上采样,这样小雨滴和大面积雨雾都能被捕获。
方向二:注意力机制。雨纹在图像中的分布是不均匀的,有些区域密集、有些区域稀疏。在残差块中加入空间注意力或者通道注意力,能帮网络自动聚焦到雨纹多的区域。实现也不复杂,常见的做法是在每层特征图后加一个 1x1 卷积生成注意力权重,把特征图和注意力权重相乘。
方向三:多阶段迭代。单次去雨往往不够彻底,尤其是大雨场景。可以设计多个阶段串联的结构,第一个阶段去掉大部分雨纹,第二个阶段针对剩余残差继续清理,类似递进式去雨。这个思路的操作方法是加载两个模型或者在一个模型里循环跑多次。
7.3 数据量不够时怎么办
深度学习项目里最尴尬的情况就是数据量不足。如果只有几十对训练图像,模型很容易过拟合到训练集上。这时候有几个补救措施。
最有效的是数据增强的加码——除了翻转裁剪,加入色彩抖动、随机旋转、高斯噪声、亮度扰动。虽然雨纹是固定的,但让背景的亮度、颜色有更多变化,模型会去学“雨纹形状”而不是“特定背景颜色”。
另外可以试试用公开的大规模数据集做预训练,再去自己的小数据集上微调。常用做法是用 Rain100H 训练 100 轮作为基础模型,然后在自己数据上用较小的学习率再训练 20 轮,通常能取得比在小数据上从头训练更好的效果。
写在实际操作之后
做完这个项目,我个人体会最深的一点是:图像去雨这个任务麻雀虽小但五脏俱全,从数据读取、数据增广、模型设计、损失函数到训练调度、指标评估,覆盖了深度学习图像处理的全部基础环节。这套代码跑完之后,你换成去雾、去噪、超分辨率,只需要改一下数据目录和模型输出通道,其他部分几乎可以原样复用。
最后再分享一个小技巧:训练过程中把每个 epoch 输出的结果图保存下来,侧边放输入图、目标图和模型输出图,做成三栏对比视觉化地看整个过程。我在训练到第 10 轮时肉眼看输出还有点模糊,到第 60 轮时雨纹基本消除,这个过程非常有成就感,也帮助你判断模型是否真的在学东西,而不是光靠 loss 数字说事。