espcn pytorch代码

ESPCN（Efficient Sub-Pixel Convolutional Neural Network）是一种用于图像超分辨率的深度学习模型，它可以将低分辨率图像转换为高分辨率图像。下面是一个简单的ESPCN的PyTorch代码示例：

import torch
import torch.nn as nn

class ESPCN(nn.Module):
    def __init__(self, scale_factor):
        super(ESPCN, self).__init__()
        self.scale_factor = scale_factor
        self.conv1 = nn.Conv2d(3, 64, kernel_size=5, padding=2)
        self.conv2 = nn.Conv2d(64, 32, kernel_size=3, padding=1)
        self.conv3 = nn.Conv2d(32, 3 * (scale_factor ** 2), kernel_size=3, padding=1)
        self.pixel_shuffle = nn.PixelShuffle(scale_factor)

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

# 创建ESPCN模型实例
scale_factor = 4  # 放大倍数
model = ESPCN(scale_factor)

# 加载预训练模型参数（可选）
# model.load_state_dict(torch.load('espcn.pth'))

# 输入低分辨率图像
input_image = torch.randn(1, 3, 32, 32)  # 输入图像尺寸为32x32

# 使用ESPCN模型进行超分辨率重建
output_image = model(input_image)

print("输入图像尺寸:", input_image.size())
print("输出图像尺寸:", output_image.size())

上述代码定义了一个名为ESPCN的PyTorch模型类，其中包含了三个卷积层和一个像素重排层。在forward方法中，输入图像经过卷积和ReLU激活函数后，通过像素重排层进行上采样，最终得到高分辨率的图像输出。

你可以根据自己的需求修改模型的结构和参数，例如调整卷积层的通道数、卷积核大小等。此外，你还可以加载预训练的模型参数（如果有的话）来进行图像超分辨率的重建。