Restormer实战：用Python从零实现图像去噪（附完整代码解析）-Seo优化-合肥市网站建设公司

Restormer实战：用Python从零实现图像去噪（附完整代码解析）

在数字图像处理领域，噪声一直是影响图像质量的关键因素。无论是医学影像、卫星遥感还是日常摄影，去除噪声都是提升图像可用性的首要步骤。传统方法如高斯滤波、中值滤波等虽然简单易用，但对于复杂噪声往往力不从心。近年来，基于深度学习的图像去噪方法展现出惊人效果，其中Restormer以其独特的Transformer架构成为业界新星。

本文将带您从零开始实现Restormer模型，不仅会详细解析核心代码，还会分享实际训练中的调参技巧。不同于单纯的理论讲解，我们更关注工程实践——如何用Python高效实现这个模型，如何处理训练数据，以及如何评估去噪效果。即使您之前没有接触过Transformer，也能通过本文的逐步指导掌握这一强大工具。

1. 环境准备与数据加载

实现Restormer的第一步是搭建合适的开发环境。推荐使用Python 3.8+和PyTorch 1.10+，这些版本在兼容性和性能上都有良好表现。以下是需要安装的核心依赖：

pip install torch torchvision torchaudio pip install opencv-python numpy tqdm matplotlib

对于图像去噪任务，数据准备至关重要。我们可以使用BSD68、DIV2K等公开数据集，也可以构建自己的噪声-干净图像对。下面是一个通用的数据加载器实现：

import torch from torch.utils.data import Dataset, DataLoader import cv2 import os import numpy as np class DenoisingDataset(Dataset): def __init__(self, clean_dir, noise_dir, transform=None): self.clean_images = [os.path.join(clean_dir, f) for f in os.listdir(clean_dir)] self.noise_images = [os.path.join(noise_dir, f) for f in os.listdir(noise_dir)] self.transform = transform def __len__(self): return min(len(self.clean_images), len(self.noise_images)) def __getitem__(self, idx): clean_img = cv2.imread(self.clean_images[idx]) noise_img = cv2.imread(self.noise_images[idx]) # 转换为PyTorch tensor并归一化 clean_img = torch.from_numpy(clean_img).float().permute(2,0,1) / 255.0 noise_img = torch.from_numpy(noise_img).float().permute(2,0,1) / 255.0 if self.transform: clean_img = self.transform(clean_img) noise_img = self.transform(noise_img) return noise_img, clean_img

提示：在实际应用中，可以考虑使用在线数据增强技术，如随机裁剪、旋转和翻转，这能有效提升模型的泛化能力。

2. Restormer核心模块实现

Restormer的核心创新在于其改进的Transformer模块，主要包括MDTA（多深度卷积转置注意力）和GDFN（门控深度前馈网络）。让我们从底层开始，逐步构建这些组件。

2.1 MDTA模块实现

MDTA是Restormer的核心注意力机制，它通过深度卷积有效降低了计算复杂度：

import torch import torch.nn as nn import torch.nn.functional as F from einops import rearrange class MDTA(nn.Module): def __init__(self, channels, num_heads=8): super(MDTA, self).__init__() self.num_heads = num_heads self.temperature = nn.Parameter(torch.ones(1, num_heads, 1, 1)) # 1x1卷积升维 self.qkv = nn.Conv2d(channels, channels*3, kernel_size=1, bias=False) # 3x3深度卷积 self.qkv_dwconv = nn.Conv2d(channels*3, channels*3, kernel_size=3, stride=1, padding=1, groups=channels*3, bias=False) # 输出投影 self.project_out = nn.Conv2d(channels, channels, kernel_size=1, bias=False) def forward(self, x): b,c,h,w = x.shape # 生成QKV qkv = self.qkv_dwconv(self.qkv(x)) q,k,v = qkv.chunk(3, dim=1) # 沿通道维度分割 # 重排为多头形式 q = rearrange(q, 'b (head c) h w -> b head c (h w)', head=self.num_heads) k = rearrange(k, 'b (head c) h w -> b head c (h w)', head=self.num_heads) v = rearrange(v, 'b (head c) h w -> b head c (h w)', head=self.num_heads) # 归一化和注意力计算 q = F.normalize(q, dim=-1) k = F.normalize(k, dim=-1) attn = (q @ k.transpose(-2, -1)) * self.temperature attn = attn.softmax(dim=-1) # 注意力加权和 out = (attn @ v) out = rearrange(out, 'b head c (h w) -> b (head c) h w', head=self.num_heads, h=h, w=w) out = self.project_out(out) return out

2.2 GDFN模块实现

GDFN模块通过门控机制控制信息流动，比传统前馈网络更高效：

class GDFN(nn.Module): def __init__(self, channels, expansion_factor=2.66): super(GDFN, self).__init__() hidden_channels = int(channels * expansion_factor) # 1x1卷积升维 self.project_in = nn.Conv2d(channels, hidden_channels*2, kernel_size=1, bias=False) # 3x3深度卷积 self.dwconv = nn.Conv2d(hidden_channels*2, hidden_channels*2, kernel_size=3, stride=1, padding=1, groups=hidden_channels*2, bias=False) # 1x1卷积降维 self.project_out = nn.Conv2d(hidden_channels, channels, kernel_size=1, bias=False) def forward(self, x): x = self.project_in(x) x1, x2 = self.dwconv(x).chunk(2, dim=1) x = F.gelu(x1) * x2 # 门控机制 x = self.project_out(x) return x

2.3 Transformer块集成

将MDTA和GDFN组合成完整的Transformer块：

class TransformerBlock(nn.Module): def __init__(self, channels, num_heads=8, expansion_factor=2.66): super(TransformerBlock, self).__init__() self.norm1 = nn.LayerNorm(channels) self.attn = MDTA(channels, num_heads) self.norm2 = nn.LayerNorm(channels) self.ffn = GDFN(channels, expansion_factor) def forward(self, x): x = x + self.attn(self.norm1(x.permute(0,2,3,1)).permute(0,3,1,2)) x = x + self.ffn(self.norm2(x.permute(0,2,3,1)).permute(0,3,1,2)) return x

3. 完整Restormer架构搭建

有了核心模块后，我们可以构建完整的Restormer模型。以下是编码器-解码器架构的实现：

class Restormer(nn.Module): def __init__(self, in_channels=3, out_channels=3, num_blocks=[4,6,6,8], num_heads=[1,2,4,8], channels=[48,96,192,384], expansion_factor=2.66): super(Restormer, self).__init__() # 下采样模块 self.down1 = nn.Sequential( nn.Conv2d(in_channels, channels[0], kernel_size=3, stride=1, padding=1), nn.Conv2d(channels[0], channels[0], kernel_size=3, stride=2, padding=1) ) self.down2 = nn.Sequential( nn.Conv2d(channels[0], channels[1], kernel_size=3, stride=1, padding=1), nn.Conv2d(channels[1], channels[1], kernel_size=3, stride=2, padding=1) ) self.down3 = nn.Sequential( nn.Conv2d(channels[1], channels[2], kernel_size=3, stride=1, padding=1), nn.Conv2d(channels[2], channels[2], kernel_size=3, stride=2, padding=1) ) # 中间Transformer块 self.transformer = nn.Sequential(*[ TransformerBlock(channels[3], num_heads[3], expansion_factor) for _ in range(num_blocks[3]) ]) # 上采样模块 self.up1 = nn.Sequential( nn.ConvTranspose2d(channels[3], channels[2], kernel_size=3, stride=2, padding=1, output_padding=1), nn.Conv2d(channels[2], channels[2], kernel_size=3, stride=1, padding=1) ) self.up2 = nn.Sequential( nn.ConvTranspose2d(channels[2], channels[1], kernel_size=3, stride=2, padding=1, output_padding=1), nn.Conv2d(channels[1], channels[1], kernel_size=3, stride=1, padding=1) ) self.up3 = nn.Sequential( nn.ConvTranspose2d(channels[1], channels[0], kernel_size=3, stride=2, padding=1, output_padding=1), nn.Conv2d(channels[0], channels[0], kernel_size=3, stride=1, padding=1) ) # 输出层 self.output = nn.Conv2d(channels[0], out_channels, kernel_size=3, stride=1, padding=1) def forward(self, x): # 编码器路径 x1 = self.down1(x) x2 = self.down2(x1) x3 = self.down3(x2) # 中间处理 x = self.transformer(x3) # 解码器路径 x = self.up1(x + x3) x = self.up2(x + x2) x = self.up3(x + x1) return self.output(x)

4. 模型训练与评估

有了模型架构后，我们需要设计合适的训练流程。以下是训练脚本的关键部分：

import torch.optim as optim from torch.utils.data import DataLoader from tqdm import tqdm def train_model(model, train_loader, val_loader, epochs=100, lr=1e-4): device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') model = model.to(device) criterion = nn.L1Loss() # 使用L1损失 optimizer = optim.AdamW(model.parameters(), lr=lr) scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, 'min', patience=5) best_psnr = 0 for epoch in range(epochs): model.train() train_loss = 0 for noisy, clean in tqdm(train_loader, desc=f'Epoch {epoch+1}'): noisy, clean = noisy.to(device), clean.to(device) optimizer.zero_grad() outputs = model(noisy) loss = criterion(outputs, clean) loss.backward() optimizer.step() train_loss += loss.item() # 验证阶段 model.eval() val_loss = 0 psnr = 0 with torch.no_grad(): for noisy, clean in val_loader: noisy, clean = noisy.to(device), clean.to(device) outputs = model(noisy) val_loss += criterion(outputs, clean).item() psnr += calculate_psnr(outputs, clean) avg_train_loss = train_loss / len(train_loader) avg_val_loss = val_loss / len(val_loader) avg_psnr = psnr / len(val_loader) print(f'Epoch {epoch+1}: Train Loss: {avg_train_loss:.4f}, Val Loss: {avg_val_loss:.4f}, PSNR: {avg_psnr:.2f}dB') # 学习率调整 scheduler.step(avg_val_loss) # 保存最佳模型 if avg_psnr > best_psnr: best_psnr = avg_psnr torch.save(model.state_dict(), 'best_restormer.pth') return model def calculate_psnr(img1, img2, max_val=1.0): mse = torch.mean((img1 - img2) ** 2) return 10 * torch.log10(max_val**2 / mse)

注意：在实际训练中，建议使用混合精度训练以节省显存并加速训练过程。可以添加以下代码：
from torch.cuda.amp import GradScaler, autocast scaler = GradScaler() # 在训练循环中替换为： with autocast(): outputs = model(noisy) loss = criterion(outputs, clean) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()

5. 实际应用与效果展示

训练完成后，我们可以使用模型对真实噪声图像进行去噪处理。以下是一个简单的推理脚本：

import cv2 import numpy as np def denoise_image(model, image_path, output_path, device='cuda'): # 加载图像 noisy_img = cv2.imread(image_path) noisy_img = cv2.cvtColor(noisy_img, cv2.COLOR_BGR2RGB) # 预处理 img_tensor = torch.from_numpy(noisy_img).float().permute(2,0,1).unsqueeze(0) / 255.0 img_tensor = img_tensor.to(device) # 推理 model.eval() with torch.no_grad(): denoised = model(img_tensor) # 后处理 denoised = denoised.squeeze().permute(1,2,0).clamp(0,1).cpu().numpy() denoised = (denoised * 255).astype(np.uint8) denoised = cv2.cvtColor(denoised, cv2.COLOR_RGB2BGR) # 保存结果 cv2.imwrite(output_path, denoised) return denoised

为了直观展示去噪效果，我们可以对比处理前后的图像质量指标：