基于pytorch实现图片去噪

以下是基于PyTorch实现图片去噪的示例代码：

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision import transforms, datasets

# 定义超参数
batch_size = 64
epochs = 10
learning_rate = 0.001

# 加载数据集
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.5,), (0.5,))
])

train_dataset = datasets.MNIST(root='./data', train=True, download=True, transform=transform)
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)

# 构建模型
class DenoiseAutoencoder(nn.Module):
    def __init__(self):
        super(DenoiseAutoencoder, self).__init__()
        self.encoder = nn.Sequential(
            nn.Linear(28*28, 256),
            nn.ReLU(),
            nn.Linear(256, 64),
            nn.ReLU(),
            nn.Linear(64, 16),
            nn.ReLU()
        )
        
        self.decoder = nn.Sequential(
            nn.Linear(16, 64),
            nn.ReLU(),
            nn.Linear(64, 256),
            nn.ReLU(),
            nn.Linear(256, 28*28),
            nn.Tanh()
        )
        
    def forward(self, x):
        x = x.view(-1, 28*28)
        x = self.encoder(x)
        x = self.decoder(x)
        x = x.view(-1, 1, 28, 28)
        return x

model = DenoiseAutoencoder()

# 定义损失函数和优化器
criterion = nn.MSELoss()
optimizer = optim.Adam(model.parameters(), lr=learning_rate)

# 训练模型
for epoch in range(epochs):
    running_loss = 0.0
    for data in train_loader:
        images, _ = data
        noisy_images = images + 0.5 * torch.randn(images.size())
        optimizer.zero_grad()
        outputs = model(noisy_images)
        loss = criterion(outputs, images)
        loss.backward()
        optimizer.step()
        running_loss += loss.item() * images.size(0)
    epoch_loss = running_loss / len(train_loader.dataset)
    print('Epoch [{}/{}], Loss: {:.4f}'.format(epoch+1, epochs, epoch_loss))

# 保存模型
torch.save(model.state_dict(), 'denoise_autoencoder.pth')

该代码使用了一个基于全连接层的自编码器来实现图片去噪。首先，我们加载MNIST数据集，并使用一个带有高斯噪声的版本来训练我们的自编码器。我们定义了一个包含三个隐藏层的自编码器，其中编码器使用ReLU激活函数，解码器使用Tanh激活函数。我们使用MSE损失函数来衡量输出图像和输入图像之间的差异，并使用Adam优化器来更新模型参数。最后，我们保存了训练好的模型。