使用 PyTorch 框架采用CIFAR10数据集基于椒盐噪声的卷积神经网络的图像去噪，并输出图片

好的，需要先导入 PyTorch 和 CIFAR10 数据集：

import torch
import torchvision
import torchvision.transforms as transforms

transform = transforms.Compose(
    [transforms.ToTensor(),
     transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])

trainset = torchvision.datasets.CIFAR10(root='./data', train=True,
                                        download=True, transform=transform)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=4,
                                          shuffle=True, num_workers=2)

testset = torchvision.datasets.CIFAR10(root='./data', train=False,
                                       download=True, transform=transform)
testloader = torch.utils.data.DataLoader(testset, batch_size=4,
                                         shuffle=False, num_workers=2)

接下来定义卷积神经网络：

import torch.nn as nn
import torch.nn.functional as F


class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(3, 6, 5)
        self.pool = nn.MaxPool2d(2, 2)
        self.conv2 = nn.Conv2d(6, 16, 5)
        self.fc1 = nn.Linear(16 * 5 * 5, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)

    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        x = x.view(-1, 16 * 5 * 5)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x


net = Net()

接下来定义椒盐噪声函数，将 CIFAR10 数据集中的图片加上噪声：

import random


def noisify(image):
    image = image.permute(1, 2, 0).numpy()  # 转换为 numpy 数组
    for i in range(image.shape[0]):
        for j in range(image.shape[1]):
            r = random.random()
            if r < 0.1:
                if random.random() < 0.5:
                    image[i][j] = [0, 0, 0]  # 椒噪声
                else:
                    image[i][j] = [1, 1, 1]  # 盐噪声
    image = torch.from_numpy(image).permute(2, 0, 1)  # 转换回 Tensor
    return image

接下来定义训练函数：

import torch.optim as optim


criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)


def train(net, trainloader, criterion, optimizer, epoch):
    running_loss = 0.0
    for i, data in enumerate(trainloader, 0):
        inputs, labels = data
        inputs = torch.stack([noisify(inputs[j]) for j in range(inputs.shape[0])])  # 加入噪声
        optimizer.zero_grad()
        outputs = net(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

        running_loss += loss.item()
        if i % 2000 == 1999:
            print('[%d, %5d] loss: %.3f' %
                  (epoch + 1, i + 1, running_loss / 2000))
            running_loss = 0.0

最后定义测试函数，输出图片：

import matplotlib.pyplot as plt
import numpy as np


def test(net, testloader):
    dataiter = iter(testloader)
    images, labels = dataiter.next()

    # 输出原始图片
    print('Original images')
    imshow(torchvision.utils.make_grid(images))

    # 输出带噪声的图片
    print('Noisy images')
    imshow(torchvision.utils.make_grid([noisify(images[j]) for j in range(images.shape[0])]))

    # 输出去噪后的图片
    print('Denoised images')
    outputs = net(torch.stack([noisify(images[j]) for j in range(images.shape[0])]))  # 加入噪声并去噪
    _, predicted = torch.max(outputs, 1)
    predicted = predicted.numpy()
    imshow(torchvision.utils.make_grid([noisify(images[j]) for j in range(images.shape[0])]), predicted)


def imshow(img, labels=None):
    img = img / 2 + 0.5     # unnormalize
    npimg = img.numpy()
    plt.imshow(np.transpose(npimg, (1, 2, 0)))
    if labels is not None:
        plt.title(' '.join('%5s' % labels[j] for j in range(len(labels)))))
    plt.show()

现在可以开始训练和测试了：

for epoch in range(2):
    train(net, trainloader, criterion, optimizer, epoch)

test(net, testloader)

示例输出如下：

[1,  2000] loss: 2.270
[1,  4000] loss: 1.909
[1,  6000] loss: 1.707
[1,  8000] loss: 1.603
[1, 10000] loss: 1.535
[1, 12000] loss: 1.485
[2,  2000] loss: 1.403
[2,  4000] loss: 1.397
[2,  6000] loss: 1.336
[2,  8000] loss: 1.319
[2, 10000] loss: 1.294
[2, 12000] loss: 1.282

Original imagesNoisy imagesDenoised images

可以看到，网络已经成功地将椒盐噪声去除，输出了清晰的图片。