使用TORCH采用CIFAR10数据集基于卷积神经网络的图像去噪 椒盐噪声

好的，以下是实现步骤：

1. 导入必要的库和模块

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

2. 定义超参数

batch_size = 128
learning_rate = 0.001
num_epochs = 50
noise_ratio = 0.2  # 椒盐噪声比例

3. 加载CIFAR10数据集，并添加椒盐噪声

train_dataset = datasets.CIFAR10(root='data', train=True, transform=transforms.ToTensor(), download=True)
test_dataset = datasets.CIFAR10(root='data', train=False, transform=transforms.ToTensor(), download=True)

# 添加椒盐噪声
def add_noise(image, noise_ratio):
    noise = np.random.rand(*image.shape)
    noisy_image = np.copy(image)
    noisy_image[noise <= (noise_ratio / 2)] = 0
    noisy_image[noise >= 1 - (noise_ratio / 2)] = 1
    return noisy_image

def add_noise_to_dataset(dataset, noise_ratio):
    noisy_dataset = []
    for image, label in dataset:
        noisy_image = add_noise(image.numpy(), noise_ratio)
        noisy_dataset.append((torch.from_numpy(noisy_image), label))
    return noisy_dataset

noisy_train_dataset = add_noise_to_dataset(train_dataset, noise_ratio)
noisy_test_dataset = add_noise_to_dataset(test_dataset, noise_ratio)

train_loader = DataLoader(dataset=noisy_train_dataset, batch_size=batch_size, shuffle=True)
test_loader = DataLoader(dataset=noisy_test_dataset, batch_size=batch_size, shuffle=False)

4. 定义卷积神经网络模型

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

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

5. 定义损失函数和优化器

model = DenoisingCNN().cuda()
criterion = nn.MSELoss()
optimizer = optim.Adam(model.parameters(), lr=learning_rate)

6. 训练模型

for epoch in range(num_epochs):
    for i, (images, labels) in enumerate(train_loader):
        images = images.cuda()
        optimizer.zero_grad()
        outputs = model(images)
        loss = criterion(outputs, images)
        loss.backward()
        optimizer.step()

        if (i+1) % 100 == 0:
            print('Epoch [{}/{}], Step [{}/{}], Loss: {:.4f}'.format(epoch+1, num_epochs, i+1, len(train_loader), loss.item()))

7. 测试模型

model.eval()
with torch.no_grad():
    total_loss = 0
    for images, labels in test_loader:
        images = images.cuda()
        outputs = model(images)
        loss = criterion(outputs, images)
        total_loss += loss * images.size(0)
    mean_loss = total_loss / len(test_loader.dataset)
    print('Test Loss: {:.4f}'.format(mean_loss.item()))

8. 可视化结果

def imshow(img, title):
    img = img.numpy().transpose((1, 2, 0))
    plt.imshow(img)
    plt.title(title)
    plt.show()

images, labels = iter(test_loader).next()
images = images.cuda()
noisy_images = images.cpu().numpy()
denoised_images = model(images).cpu().numpy()

for i in range(5):
    plt.figure(figsize=(10, 4))
    plt.subplot(131)
    imshow(noisy_images[i], 'Noisy Image')
    plt.subplot(132)
    imshow(denoised_images[i], 'Denoised Image')
    plt.subplot(133)
    imshow(images[i].cpu(), 'Original Image')

这样就完成了基于卷积神经网络的图像去噪（椒盐噪声）任务。