简介：本资源是一套面向图像处理研究者与深度学习开发者的优质图像去噪实战项目，聚焦于解决真实场景中高斯噪声、泊松噪声等常见退化问题，适用于医学影像增强、卫星图像复原及低光照摄影后处理等实际应用。项目创新性融合Swin-Transformer的长程建模能力与UNet的多尺度特征融合机制，构建端到端去噪网络SUNet，并配套设计专用损失函数以兼顾结构保真与视觉自然性。压缩包共25个文件，含17个核心Python脚本（如train.py、demo.py、SUNet.py、data_RGB.py）、3个MATLAB评估脚本（含DIV2K噪声数据生成与指标计算）、2份Markdown文档（README说明与使用指南）及1个训练配置yaml文件，整体仅37KB，轻量易部署。已有576人学习下载，提供从数据预处理、模型训练、任意分辨率推理到定量评估的完整闭环实现，代码结构清晰、模块解耦良好，支持快速复现与二次开发。
1. 项目概述：当Swin-Transformer遇上UNet，图像去噪的新解法
最近在整理手头的图像处理项目时，翻出了一个让我印象深刻的实战案例：一个结合了Swin-Transformer和UNet架构的图像去噪模型。这个项目最初是为了解决一个实际问题——处理一批在低光照条件下拍摄、带有复杂噪声的医学显微图像。传统的滤波方法效果有限，而当时（现在依然是）主流的基于CNN的UNet虽然强大，但在捕捉长距离依赖和全局上下文信息上总觉得差那么点意思。于是，我们尝试将视觉Transformer领域的新星Swin-Transformer引入到UNet的编码器部分，结果出乎意料地好，不仅在公开数据集上刷出了不错的分数，在实际应用中也表现稳定。这个项目打包了完整的训练、推理代码和预训练模型，对于想深入理解现代图像复原技术，特别是想动手实践Transformer与CNN混合架构的朋友来说，是个非常不错的切入点。无论你是计算机视觉的研究者，还是希望将先进算法落地的工程师，都能从中获得直接的代码参考和设计思路。
2. 核心架构解析：为什么是Swin-Transformer + UNet？
2.1 UNet的基石作用与固有局限
UNet结构在图像分割领域一战成名，其经典的“编码器-解码器”加“跳跃连接”的设计，使其在图像到图像的任务中几乎成为默认的基线模型。编码器负责下采样，提取多层次、由浅入深的特征；解码器负责上采样，逐步恢复空间分辨率并融合深层语义信息；跳跃连接则将编码器每一层的特征图直接传递到解码器对应层，有效缓解了梯度消失问题，并保留了更多的空间细节信息。对于图像去噪这类像素级回归任务，UNet这种能同时兼顾全局上下文和局部细节的结构，天然具有优势。
然而，传统UNet的编码器通常由堆叠的卷积层和池化层构成。卷积操作的感受野是局部的，尽管通过堆叠可以扩大，但其捕捉长距离像素间关系（例如图像中两个相隔很远的相似纹理区域之间的相关性）的能力是间接且效率较低的。而图像噪声，尤其是结构化噪声或与内容相关的复杂噪声，其模式可能遍布整个图像，需要模型具备强大的全局建模能力才能有效区分噪声与真实信号。
2.2 Swin-Transformer的引入：补全全局建模短板
这正是Swin-Transformer大显身手的地方。Swin-Transformer的核心创新在于其“分层设计”和“移位窗口”机制。它将图像分割成一个个不重叠的局部窗口，在每个窗口内计算自注意力（Window Multi-Head Self-Attention, W-MSA）。这种设计将计算复杂度从图像尺寸的平方级降低到线性级，使得处理高分辨率图像成为可能。更重要的是，通过下一层窗口的偏移（Shifted Windows），不同窗口之间的信息得以交互，从而在有限的层数内建立起整个图像的全局依赖关系。
将Swin-Transformer模块作为UNet编码器的核心构建块，意味着我们在特征提取的每一个阶段，都能在局部窗口内进行高效的全局交互（窗口内自注意力），并通过层级结构和移位窗口机制，逐步建立起跨越大区域的语义关联。这相当于给UNet装上了一双“全局眼”，使其不仅能看清像素周围的细节（靠卷积或窗口内的自注意力），还能理解图像中遥远部分之间的关联（靠移位窗口和层级传递），这对于区分广泛分布的噪声模式和真实的图像结构至关重要。
2.3 混合架构的设计哲学与具体实现
在我们的项目中，没有完全用Transformer块替换所有卷积。一个典型的混合设计是：使用Swin-Transformer块构建编码器的主干，而在解码器部分以及跳跃连接后的特征融合处，仍然使用卷积层。这样设计的考量是：
- 发挥各自专长：Swin-Transformer擅长捕获全局和长程依赖，适合在编码阶段学习鲁棒的、富含语义的特征表示。而卷积层具有固有的平移等变性和局部性，在解码阶段进行上采样和精细的像素级重建时更加高效和稳定。
- 降低计算成本：将计算量较大的Transformer模块主要放在下采样路径上，随着特征图尺寸减小，其计算负担也随之降低。上采样路径保持轻量的卷积操作，有利于快速推理。
- 保持细节信息：跳跃连接传递的是来自Swin-Transformer编码器各阶段的特征。这些特征已经融入了全局信息，再与解码器的卷积特征图结合，能确保最终输出的去噪图像既具有全局一致性，又不丢失重要的纹理细节。
具体到网络结构，输入噪声图像首先经过一个卷积层进行浅层特征提取。然后进入四个下采样阶段，每个阶段由若干个Swin-Transformer块和一个Patch Merging层（实现下采样）组成。解码器对应四个上采样阶段，每个阶段通过转置卷积或像素洗牌操作上采样特征图，并与来自编码器的对应特征图拼接（跳跃连接），再经过卷积块进行特征融合。最后，一个卷积层将通道数映射回3（RGB）或1（灰度），输出预测的干净图像。
注意：这里的一个关键细节是位置编码。由于Swin-Transformer对输入序列的顺序敏感，我们需要为图像块添加可学习的位置编码。在Swin-Transformer的原论文中，使用了相对位置偏置（relative position bias），这对于处理可变尺寸输入和获得平移不变性的泛化能力很有帮助。在我们的实现中，直接采用了这一设计。
3. 项目实战：从环境搭建到模型训练
3.1 开发环境与依赖配置
这个项目基于PyTorch框架。首先确保你的机器上安装了合适版本的CUDA和cuDNN（如果你有NVIDIA GPU的话），这对于训练Transformer模型是几乎必须的，能极大加速训练过程。
# 推荐使用conda创建虚拟环境 conda create -n image-denoising python=3.8 conda activate image-denoising # 安装PyTorch (请根据你的CUDA版本访问PyTorch官网获取对应命令) # 例如，对于CUDA 11.3 pip install torch==1.12.1+cu113 torchvision==0.13.1+cu113 torchaudio==0.12.1 --extra-index-url https://download.pytorch.org/whl/cu113 # 安装项目核心依赖 pip install opencv-python pillow matplotlib scikit-image tqdm tensorboard # 安装Swin-Transformer的PyTorch实现，例如 timm 库 pip install timm # 或者，如果你直接使用项目源码中的模型定义，则无需单独安装timm项目源码结构通常清晰明了：
config/: 存放模型和训练参数的配置文件（如config.yaml）。data/: 数据加载和预处理模块。models/: 网络模型定义，核心的swin_unet.py就在这里。utils/: 工具函数，如损失函数、指标计算、日志记录等。train.py: 模型训练主脚本。test.py/inference.py: 模型测试和推理脚本。pretrained_models/: 存放预训练权重（如果有的话）。
3.2 数据准备与预处理策略
高质量的数据是模型成功的基石。对于图像去噪，你需要准备“噪声-干净”图像对。常见的数据集有：
- 合成数据集：如使用BSD68、Set12等干净数据集，人工添加高斯噪声、泊松噪声或更加复杂的盲噪声。优点是噪声水平可控，便于定量评估。
- 真实噪声数据集：如SIDD、DND。这些数据集包含了真实场景下拍摄的噪声图像及其对应的（多帧平均得到的）近似干净图像。挑战更大，也更贴近实际应用。
在我们的项目中，为了快速验证和演示，通常先从合成高斯噪声开始。数据预处理管道包括：
- 读取与归一化：将图像像素值从[0, 255]归一化到[0, 1]或[-1, 1]。PyTorch的
ToTensor()变换会自动除以255。 - 添加噪声：在训练时动态添加噪声，可以增加数据的多样性。例如，对于高斯噪声，噪声水平σ可以在一个范围内随机采样。
import torch def add_gaussian_noise(clean_img, noise_sigma_range=(0, 50)): sigma = torch.rand(1) * (noise_sigma_range[1] - noise_sigma_range[0]) / 255.0 + noise_sigma_range[0] / 255.0 noise = torch.randn_like(clean_img) * sigma noisy_img = clean_img + noise # 确保像素值在合法范围内 noisy_img = torch.clamp(noisy_img, 0., 1.) return noisy_img, clean_img - 数据增强：虽然去噪任务对几何变换敏感（因为需要像素级对齐），但我们可以使用随机水平/垂直翻转、旋转90度等增强方式。切记，必须对“噪声-干净”图像对施加完全相同的变换，以保证它们的空间对齐。
- Patch化：对于高分辨率图像，直接输入网络可能受限于GPU内存。常见的做法是随机裁剪出固定大小的图像块（如128x128, 256x256）进行训练。这同时也是一种数据增强。
3.3 模型训练的关键步骤与超参数调优
训练脚本train.py是项目的核心。其主要流程如下：
- 初始化：解析配置文件，设置随机种子（保证可复现性），创建数据加载器、模型、优化器、损失函数和评估指标。
- 损失函数选择：图像去噪最常用的损失是L1损失（Mean Absolute Error, MAE）和L2损失（Mean Squared Error, MSE，即PSNR的优化目标）。L1损失对异常值不那么敏感，训练出的图像边缘可能更清晰；L2损失更注重整体像素误差的平方和，通常能获得更高的PSNR。实践中，可以结合使用，例如
Loss = α * L1 + β * L2。我们还尝试了感知损失（Perceptual Loss），即利用预训练网络（如VGG）提取特征图计算差异，以提升视觉质量，但这会显著增加计算量。 - 优化器与学习率调度：Adam或AdamW优化器是当前的首选。对于Swin-Transformer，AdamW（带权重衰减的Adam）效果通常更好，能有效防止过拟合。学习率采用带热启动（Warmup）的余弦退火（Cosine Annealing）策略是Transformer训练的标配。Warmup让模型在训练初期以较小的学习率“热身”，避免梯度不稳定；余弦退火则在训练中后期平滑地降低学习率，有助于模型收敛到更优的局部最优点。
# 配置文件 config.yaml 片段示例 optimizer: name: 'AdamW' lr: 1e-4 weight_decay: 0.05 scheduler: name: 'CosineAnnealingLR' T_max: 300 # 总epoch数 eta_min: 1e-6 warmup: epochs: 5 start_lr: 1e-7 - 训练循环：标准的PyTorch训练循环。在每个epoch中，遍历训练集，前向传播计算损失，反向传播更新权重。定期在验证集上评估模型性能（如PSNR, SSIM），并保存验证集上表现最好的模型权重。
- 监控与可视化：使用TensorBoard或WandB记录训练损失、验证指标、学习率变化，并可视化一些样例输入-输出对。这对于调试和选择最佳模型至关重要。
实操心得：训练Swin-Transformer+UNet这类模型，显存消耗较大。如果遇到CUDA out of memory错误，首先尝试减小批量大小（batch size）或输入图像块的大小。其次，可以使用混合精度训练（AMP），这能显著减少显存占用并可能加快训练速度，且通常不会影响最终精度。在PyTorch中，只需几行代码即可启用。
4. 核心代码模块深度解读
4.1 Swin-Transformer Block的集成
项目源码中最关键的部分是如何将Swin-Transformer块嵌入到UNet中。我们通常不会从头实现Swin-Transformer，而是借鉴或直接使用timm库中现成的、经过良好测试的模块。
import torch import torch.nn as nn import torch.nn.functional as F from timm.models.swin_transformer import SwinTransformerBlock class SwinTransformerStage(nn.Module): """ 一个Swin-Transformer阶段，包含多个连续的SwinTransformerBlock。 """ def __init__(self, dim, input_resolution, depth, num_heads, window_size, mlp_ratio=4.): super().__init__() self.blocks = nn.ModuleList() for i in range(depth): self.blocks.append( SwinTransformerBlock( dim=dim, input_resolution=input_resolution, num_heads=num_heads, window_size=window_size, shift_size=0 if (i % 2 == 0) else window_size // 2, # 交替使用常规窗口和移位窗口 mlp_ratio=mlp_ratio, ) ) def forward(self, x): for blk in self.blocks: x = blk(x) return x # 在UNet编码器中的使用示例 class EncoderBlock(nn.Module): def __init__(self, in_channels, out_channels, input_resolution, depth, num_heads, window_size): super().__init__() self.downsample = nn.Conv2d(in_channels, out_channels, kernel_size=2, stride=2) # Patch Merging的简化 # 调整分辨率以适应Transformer Block self.input_resolution = (input_resolution[0]//2, input_resolution[1]//2) # 将特征图展平为序列 (B, C, H, W) -> (B, H*W, C) self.norm = nn.LayerNorm(out_channels) self.swin_stage = SwinTransformerStage( dim=out_channels, input_resolution=self.input_resolution, depth=depth, num_heads=num_heads, window_size=window_size, ) # 将序列恢复为特征图 (B, H*W, C) -> (B, C, H, W) def forward(self, x): x = self.downsample(x) B, C, H, W = x.shape x_flat = x.flatten(2).transpose(1, 2) # (B, C, H, W) -> (B, H*W, C) x_flat = self.norm(x_flat) x_flat = self.swin_stage(x_flat) x = x_flat.transpose(1, 2).view(B, C, H, W) # (B, H*W, C) -> (B, C, H, W) return x关键点解析：
- 维度转换：SwinTransformerBlock的输入输出是序列格式
(Batch, Num_Patches, Channel)。因此，在卷积特征图输入前，需要将其从(B, C, H, W)展平并转置。处理完后，再变换回来。 - 移位窗口：通过
shift_size参数控制，每隔一个块使用移位窗口，实现跨窗口通信。 - Patch Merging：原版Swin-Transformer使用专门的Patch Merging层进行下采样。在简化实现中，我们有时会用步长为2的卷积代替，但更严谨的做法是实现其细节，包括线性变换和降维。
4.2 跳跃连接与特征融合
UNet的灵魂在于跳跃连接。在我们的混合架构中，跳跃连接传递的是经过Swin-Transformer编码器层处理后的、富含全局信息的特征图。
class DecoderBlock(nn.Module): def __init__(self, in_channels, skip_channels, out_channels): super().__init__() # 上采样：可以使用转置卷积、双线性插值+卷积或像素洗牌 self.up = nn.ConvTranspose2d(in_channels, in_channels // 2, kernel_size=2, stride=2) # 拼接跳跃连接的特征 self.conv1 = nn.Conv2d(in_channels // 2 + skip_channels, out_channels, kernel_size=3, padding=1) self.bn1 = nn.BatchNorm2d(out_channels) self.relu = nn.ReLU(inplace=True) self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1) self.bn2 = nn.BatchNorm2d(out_channels) def forward(self, x, skip): # x: 来自解码器上一层的特征 # skip: 来自编码器的对应层特征（跳跃连接） x = self.up(x) # 确保skip和x的空间尺寸一致（由于下采样/上采样可能因奇数尺寸产生1像素误差） diffY = skip.size()[2] - x.size()[2] diffX = skip.size()[3] - x.size()[3] x = F.pad(x, [diffX // 2, diffX - diffX // 2, diffY // 2, diffY - diffY // 2]) # 拼接操作 x = torch.cat([x, skip], dim=1) x = self.relu(self.bn1(self.conv1(x))) x = self.relu(self.bn2(self.conv2(x))) return x注意事项：由于Swin-Transformer块内部可能包含LayerNorm，而跳跃连接过来的特征图在编码器路径中可能已经过归一化。在解码器进行特征融合时，要小心处理特征分布的差异。通常，在拼接后立即接一个BatchNorm层或InstanceNorm层有助于稳定训练。
4.3 损失函数与评估指标的实现
除了标准的L1/L2损失，实现结构相似性指数（SSIM）作为损失或评估指标能更好地对齐人眼感知。
import torch from torch import nn from pytorch_msssim import ssim, ms_ssim # 可以使用第三方库，或自己实现 class HybridLoss(nn.Module): def __init__(self, alpha=0.84, l1_weight=1.0, ssim_weight=0.16): super().__init__() self.l1_loss = nn.L1Loss() self.alpha = alpha # 用于SSIM计算 self.l1_weight = l1_weight self.ssim_weight = ssim_weight def forward(self, pred, target): l1 = self.l1_loss(pred, target) # 计算SSIM， data_range通常为1（归一化后）或255 ssim_loss = 1 - ssim(pred, target, data_range=1.0, size_average=True) # 组合损失 total_loss = self.l1_weight * l1 + self.ssim_weight * ssim_loss return total_loss, {'L1': l1.item(), 'SSIM': 1-ssim_loss.item()} # 评估指标：PSNR def calculate_psnr(img1, img2, data_range=1.0): mse = torch.mean((img1 - img2) ** 2) if mse == 0: return float('inf') return 20 * torch.log10(data_range / torch.sqrt(mse))为什么用L1+SSIM？L1损失能有效减少模糊，保持边缘；SSIM损失则关注结构相似性，能提升视觉效果的整体自然度。两者结合在主观质量上往往优于单一的L2损失。
5. 训练技巧与调优实战经验
5.1 学习率策略与优化器设置
Transformer模型对优化策略非常敏感。我们的经验是：
- 使用AdamW而非Adam：权重衰减（Weight Decay）对于Transformer防止过拟合至关重要。AdamW将权重衰减与梯度更新解耦，效果更佳。权重衰减系数通常设置在
0.05左右。 - 必不可少的热启动（Warmup）：在训练最初的5-10个epoch，将学习率从一个小值（如
1e-7）线性增长到初始学习率（如1e-4）。这给了模型参数一个稳定的初始化适应期。 - 余弦退火（Cosine Annealing）：将学习率按照余弦函数从初始值衰减到一个非常小的最小值（如
1e-6）。T_max设置为总的训练epoch数。 - 梯度裁剪（Gradient Clipping）：虽然不像RNN中那么必要，但在训练深度Transformer时，设置一个梯度范数阈值（如
max_norm=1.0）可以避免训练初期因异常样本导致的梯度爆炸，增加训练稳定性。
5.2 数据增强的针对性设计
对于去噪任务，标准的数据增强需要谨慎：
- 几何变换：随机水平/垂直翻转、90度旋转是安全的，因为它们保持了像素间的相对位置关系。
- 色彩抖动/亮度对比度调整：谨慎使用。因为噪声特性可能与原始图像信号相关（例如，泊松噪声与亮度相关）。改变图像亮度可能会改变噪声模型，导致模型学习到错误的映射。对于合成已知噪声的任务，最好避免；对于真实盲去噪，轻微的调整或许能增加鲁棒性，但需要实验验证。
- CutMix/MixUp：这类混合样本的数据增强在分类任务中很有效，但在像素级回归任务中会破坏“噪声-干净”对的严格对应关系，通常不适用。
- 更有效的增强：在噪声层面做文章。例如，在训练时动态变化噪声水平（σ）、噪声类型（混合高斯和泊松噪声），甚至使用更复杂的噪声模型来合成数据，这能极大地提升模型对不同噪声的泛化能力，即所谓的“盲去噪”能力。
5.3 模型初始化与预训练权重利用
- 卷积层初始化：使用He初始化（Kaiming初始化）配合ReLU激活函数。
- Transformer层初始化：Swin-Transformer块中的线性层、LayerNorm层通常有默认的初始化方式，遵循原论文即可。一个技巧是，将Transformer块最后一层的线性投影权重初始化为零，这可以在训练开始时使残差分支为零，让模型更关注恒等映射，有助于稳定深度网络的训练。
- 加载预训练权重：这是大幅提升性能和收敛速度的秘诀。Swin-Transformer通常在ImageNet这样的大型数据集上进行预训练。我们可以加载其在ImageNet上预训练好的编码器权重（
timm库提供了这些权重），来初始化我们UNet编码器中的Swin-Transformer部分。这相当于让模型从一个非常好的视觉特征提取器开始学习，而不是从零开始。对于解码器和跳跃连接处的卷积层，可以随机初始化。这种“编码器预训练，解码器微调”的策略非常有效。
6. 推理部署与性能优化
6.1 模型推理与后处理
训练完成后，使用test.py或inference.py脚本进行推理。流程很简单：加载模型权重，读取噪声图像，归一化，前向传播，反归一化，保存结果。
def denoise_single_image(model, noisy_img_path, device='cuda'): model.eval() # 切换到评估模式 with torch.no_grad(): # 1. 读取并预处理图像 noisy_img = cv2.imread(noisy_img_path) noisy_img = cv2.cvtColor(noisy_img, cv2.COLOR_BGR2RGB) / 255.0 input_tensor = torch.from_numpy(noisy_img).float().permute(2,0,1).unsqueeze(0).to(device) # 2. 模型推理 output_tensor = model(input_tensor) # 3. 后处理 denoised_img = output_tensor.squeeze().cpu().permute(1,2,0).numpy() denoised_img = np.clip(denoised_img * 255, 0, 255).astype(np.uint8) denoised_img = cv2.cvtColor(denoised_img, cv2.COLOR_RGB2BGR) return denoised_img关于滑动窗口推理：对于远大于训练时裁剪尺寸的高分辨率图像，直接输入网络可能不行。常用的方法是使用滑动窗口（Overlapping Tiles）策略：将大图切割成重叠的小块，分别去噪，然后融合。融合时，对重叠区域使用加权平均（如余弦窗），以消除块边界可能产生的接缝。
6.2 模型轻量化与加速探索
Swin-Transformer+UNet模型参数量和计算量相对较大。在实际部署，特别是端侧部署时，需要考虑优化：
- 知识蒸馏：训练一个更小的学生网络（如纯卷积的轻量UNet），用训练好的大模型（教师网络）的输出作为软标签来指导学生网络训练，以期用小模型获得接近大模型的性能。
- 模型剪枝：识别并剪枝网络中不重要的连接或通道。对于Transformer，可以尝试剪枝注意力头或MLP中间层的维度。
- 量化：将模型权重和激活从32位浮点数（FP32）转换为8位整数（INT8）。PyTorch提供了方便的量化工具。量化后模型大小减少约75%，推理速度提升，但可能会带来轻微的精度损失。
- 使用更高效的架构：可以考虑用MobileViT、EfficientFormer等轻量级Vision Transformer替代Swin-Transformer，或者探索动态稀疏注意力机制来减少计算量。
6.3 实际应用中的挑战与应对
将模型应用于真实场景，会碰到训练时未曾遇到的问题：
- 领域差距：在合成高斯噪声上训练的模型，处理手机拍摄的真实噪声时效果可能下降。解决方案是进行领域适应，即在少量真实噪声图像对上继续微调模型。
- 噪声类型未知：实际噪声可能是高斯、泊松、椒盐噪声的混合，且水平未知。这就需要训练一个盲去噪模型。一种方法是使用“噪声到噪声”的学习方法，即仅用噪声图像对进行训练，或者使用更复杂的噪声建模和生成方式。
- 计算资源限制：在边缘设备上运行大模型困难。除了上述轻量化方法，还可以考虑模型分片或使用神经架构搜索自动设计适合目标平台的轻量网络。
这个Swin-Transformer+UNet的图像去噪项目，从一个具体的需求出发，融合了CNN的局部细节建模能力和Transformer的全局上下文理解能力，代表了一种当前主流的视觉任务架构设计思路。通过动手实践这个项目，你不仅能获得一个效果不错的去噪工具，更能深入理解如何将前沿的学术成果转化为解决实际问题的工程方案。项目源码中的每一个模块、每一处设计选择，都值得细细揣摩和尝试修改，这才是从“会用”到“理解”再到“创新”的关键一步。
本文还有配套的精品资源，点击获取
