PyTorch 的超分辨率重建代码

以下是一个简单的 PyTorch 超分辨率重建代码示例：

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

# 定义超分辨率模型
class SuperResolutionModel(nn.Module):
    def __init__(self, upscale_factor):
        super(SuperResolutionModel, self).__init__()
        self.upscale_factor = upscale_factor
        self.conv1 = nn.Conv2d(in_channels=3, out_channels=64, kernel_size=5, padding=2)
        self.conv2 = nn.Conv2d(in_channels=64, out_channels=32, kernel_size=3, padding=1)
        self.conv3 = nn.Conv2d(in_channels=32, out_channels=3 * (upscale_factor ** 2), kernel_size=3, padding=1)
        self.pixel_shuffle = nn.PixelShuffle(upscale_factor)
        self.relu = nn.ReLU()

    def forward(self, x):
        x = self.relu(self.conv1(x))
        x = self.relu(self.conv2(x))
        x = self.pixel_shuffle(self.conv3(x))
        return x

# 加载数据集
train_dataset = CIFAR10(root='data/', train=True, download=True, transform=ToTensor())
train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)

# 实例化模型和优化器
model = SuperResolutionModel(upscale_factor=2)
optimizer = optim.Adam(model.parameters(), lr=0.001)

# 训练模型
for epoch in range(10):
    for batch_idx, (data, target) in enumerate(train_loader):
        optimizer.zero_grad()
        output = model(data)
        loss = nn.MSELoss()(output, target)
        loss.backward()
        optimizer.step()
        if batch_idx % 100 == 0:
            print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
                epoch, batch_idx * len(data), len(train_loader.dataset),
                100. * batch_idx / len(train_loader), loss.item()))

在这个示例中，我们定义了一个简单的超分辨率模型 `SuperResolutionModel`，它包含三个卷积层和一个像素重排层。我们使用 CIFAR10 数据集进行训练，优化器使用 Adam，损失函数为均方误差（MSE）。在训练过程中，我们每 100 个 batch 打印一次损失值。