卷积神经网络图像去噪 代码

以下是一个简单的卷积神经网络图像去噪的代码示例。这个代码使用的是 PyTorch 框架。

首先需要导入必要的库：

import torch
import torch.nn as nn
import torch.optim as optim
import torchvision.transforms as transforms
from torch.utils.data import DataLoader
from torchvision.datasets import MNIST
import numpy as np
import matplotlib.pyplot as plt

接下来定义一个简单的卷积神经网络模型：

class DenoisingAutoencoder(nn.Module):
    def __init__(self):
        super(DenoisingAutoencoder, self).__init__()
        self.encoder = nn.Sequential(
            nn.Conv2d(1, 16, kernel_size=3, stride=1, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2),
            nn.Conv2d(16, 8, kernel_size=3, stride=1, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2),
            nn.Conv2d(8, 8, kernel_size=3, stride=1, padding=1),
            nn.ReLU()
        )
        self.decoder = nn.Sequential(
            nn.ConvTranspose2d(8, 8, kernel_size=3, stride=1, padding=1),
            nn.ReLU(),
            nn.ConvTranspose2d(8, 16, kernel_size=3, stride=2, padding=1, output_padding=1),
            nn.ReLU(),
            nn.ConvTranspose2d(16, 1, kernel_size=3, stride=2, padding=1, output_padding=1),
            nn.Sigmoid()
        )

    def forward(self, x):
        encoded = self.encoder(x)
        decoded = self.decoder(encoded)
        return decoded

这个模型有两个部分：编码器和解码器。编码器将输入图像压缩成一个低维度的向量，解码器将这个向量还原成原始图像。这里使用了卷积层和池化层来提取图像特征。

接下来定义训练过程：

def train(model, train_loader, num_epochs=10):
    criterion = nn.MSELoss()
    optimizer = optim.Adam(model.parameters(), lr=0.001)

    for epoch in range(num_epochs):
        running_loss = 0.0
        for data in train_loader:
            img, _ = data
            noisy_img = add_noise(img)
            optimizer.zero_grad()
            outputs = model(noisy_img)
            loss = criterion(outputs, img)
            loss.backward()
            optimizer.step()
            running_loss += loss.item()
        print('Epoch [%d], Loss: %.4f' % (epoch+1, running_loss/len(train_loader)))

这个函数将模型和训练数据作为参数，通过多次迭代来训练模型。在每个 epoch 中，它将遍历整个数据集，并计算模型的损失。这里使用了均方误差损失函数和 Adam 优化器。

最后定义一个添加噪声的函数：

def add_noise(img):
    noise = torch.randn(img.size()) * 0.2
    noisy_img = img + noise
    return noisy_img

这个函数将输入图像添加高斯噪声。这里使用了一个标准差为 0.2 的高斯分布。

现在我们可以加载数据并开始训练模型了：

transform = transforms.Compose([
    transforms.ToTensor()
])

train_set = MNIST(root='data', train=True, transform=transform, download=True)
train_loader = DataLoader(train_set, batch_size=128, shuffle=True)

model = DenoisingAutoencoder()
train(model, train_loader, num_epochs=10)

这里使用了 MNIST 数据集，但只使用了图像数据，不使用标签数据。训练过程可能需要一些时间，具体取决于你的硬件和数据集的大小。一旦训练完成，你可以使用模型来去除新图像中的噪声：

test_set = MNIST(root='data', train=False, transform=transform, download=True)
test_loader = DataLoader(test_set, batch_size=10, shuffle=True)

with torch.no_grad():
    for data in test_loader:
        img, _ = data
        noisy_img = add_noise(img)
        outputs = model(noisy_img)
        plt.figure(figsize=(20,4))
        for i in range(10):
            plt.subplot(2, 10, i+1)
            plt.imshow(noisy_img[i][0], cmap='gray')
            plt.subplot(2, 10, i+11)
            plt.imshow(outputs[i][0], cmap='gray')
        plt.show()

这个代码块将加载测试数据集，然后将一批噪声图像和去噪后的图像可视化。