CAE卷积自编码器实战：从零构建图像降噪模型

1次阅读

没有评论

共计 2399 个字符，预计需要花费 6 分钟才能阅读完成。

刚接触图像处理时，我最先学会的是用 OpenCV 的 cv2.GaussianBlur() 做降噪。但实际使用时发现，这种传统方法有三个致命伤：

细节丢失：高斯滤波在抹除噪声的同时，会把边缘信息也模糊掉
参数敏感：kernel 大小和 sigma 值需要反复手动调整
通用性差：对椒盐噪声效果尚可，但面对 CT 扫描的散斑噪声就束手无策

而卷积自编码器 (CAE) 的无监督学习特性，让它能从数据本身学习噪声分布，实现更智能的降噪。

传统自编码器 (AE) 用全连接层构建，就像用渔网捞鱼：

输入图像被展平成一维向量（丢失空间信息）
隐层节点与所有像素连接（参数爆炸）
重建图像容易出现模糊块效应

CAE 则像用带网兜的鱼叉：

卷积核只在局部感受野滑动（保留空间结构）
权值共享大幅减少参数量
MaxPooling 记录特征位置，上采样恢复细节

先安装必要的库，建议使用 Python 3.8+ 和 PyTorch 1.10+：

import torch
import torch.nn as nn
import torchvision.transforms as transforms
from torchvision.datasets import CIFAR10

下采样部分采用经典 ”Conv+ReLU+Pool” 组合，注意每层的输出尺寸变化：

class Encoder(nn.Module):
    def __init__(self):
        super().__init__()
        # 第一层：3 通道输入，16 个 3x3 卷积核，边缘填充 1 保持尺寸
        self.conv1 = nn.Conv2d(3, 16, kernel_size=3, stride=1, padding=1)
        self.relu1 = nn.ReLU()
        # 2x2 最大池化，尺寸减半
        self.pool1 = nn.MaxPool2d(2)

        # 第二层同理，通道数翻倍
        self.conv2 = nn.Conv2d(16, 32, 3, 1, 1)
        self.relu2 = nn.ReLU()
        self.pool2 = nn.MaxPool2d(2)

上采样使用转置卷积(Transpose Conv)，可以理解为反向的卷积操作：

class Decoder(nn.Module):
    def __init__(self):
        super().__init__()
        # 转置卷积：用 stride= 2 实现 2 倍上采样
        self.deconv1 = nn.ConvTranspose2d(32, 16, 
                                         kernel_size=3, 
                                         stride=2, 
                                         padding=1,
                                         output_padding=1)
        self.relu1 = nn.ReLU()

        self.deconv2 = nn.ConvTranspose2d(16, 3, 
                                         kernel_size=3,
                                         stride=2,
                                         padding=1,
                                         output_padding=1)
        self.sigmoid = nn.Sigmoid()  # 输出压缩到 [0,1] 范围

在 CIFAR-10 上模拟真实噪声场景：

# 定义噪声添加函数
def add_gaussian_noise(img, mean=0, std=0.1):
    noise = torch.randn(img.size()) * std + mean
    noisy_img = img + noise
    return torch.clamp(noisy_img, 0, 1)  # 限制在有效像素范围

# 数据加载时实时加噪声
train_transform = transforms.Compose([transforms.ToTensor(),
    transforms.Lambda(lambda x: add_gaussian_noise(x))
])

经过多次实验，推荐以下配置：

学习率：0.001（Adam 优化器）
Batch Size：根据显存选择 32/64
损失函数：MSELoss（均方误差）
Epoch 数：50 左右开始观察损失曲线

当训练数据不足时（如医疗影像）：

使用 torchvision.transforms 做弹性变换
尝试 MixUp 数据增强：混合两张图片生成新样本
在 Encoder 后添加 Dropout 层（rate=0.3）防止过拟合

遇到 CUDA out of memory 错误时：

逐步减小 batch_size（从 64→32→16）
使用 torch.cuda.empty_cache() 手动释放缓存
考虑梯度累积：每 4 个 mini-batch 更新一次参数

训练完成后，对比噪声图像和重建结果：

import matplotlib.pyplot as plt

# 随机选取测试样本
with torch.no_grad():
    noisy, clean = test_dataset[0]
    reconstructed = model(noisy.unsqueeze(0))

plt.figure(figsize=(10,4))
plt.subplot(1,3,1).set_title('Noisy Input')
plt.imshow(noisy.permute(1,2,0))
plt.subplot(1,3,2).set_title('CAE Output')
plt.imshow(reconstructed.squeeze().permute(1,2,0))
plt.subplot(1,3,3).set_title('Ground Truth')
plt.imshow(clean.permute(1,2,0))