位置：首页-资讯-后端开发

Pytorch怎么搭建SRGAN平台提升图片超分辨率

2023-06-30 11:46

短信预约 -IT技能 免费直播动态提醒

本篇内容介绍了“Pytorch怎么搭建SRGAN平台提升图片超分辨率”的有关知识，在实际案例的操作过程中，不少人都会遇到这样的困境，接下来就让小编带领大家学习一下如何处理这些情况吧！希望大家仔细阅读，能够学有所成！

网络构建

一、什么是SRGAN

SRGAN出自论文Photo-Realistic Single Image Super-Resolution Using a Generative Adversarial Network。

如果将SRGAN看作一个黑匣子，其主要的功能就是输入一张低分辨率图片，生成高分辨率图片。

Pytorch怎么搭建SRGAN平台提升图片超分辨率

该文章提到，普通的超分辨率模型训练网络时只用到了均方差作为损失函数，虽然能够获得很高的峰值信噪比，但是恢复出来的图像通常会丢失高频细节。

SRGAN利用感知损失(perceptual loss)和对抗损失(adversarial loss)来提升恢复出的图片的真实感。

二、生成网络的构建

Pytorch怎么搭建SRGAN平台提升图片超分辨率

生成网络的构成如上图所示，生成网络的作用是输入一张低分辨率图片，生成高分辨率图片。：

SRGAN的生成网络由三个部分组成。

低分辨率图像进入后会经过一个卷积+RELU函数。

然后经过B个残差网络结构，每个残差结构都包含两个卷积+标准化+RELU，还有一个残差边。

然后进入上采样部分，在经过两次上采样后，原图的高宽变为原来的4倍，实现分辨率的提升。

前两个部分用于特征提取，第三部分用于提高分辨率。

import mathimport torchfrom torch import nnclass ResidualBlock(nn.Module):    def __init__(self, channels):        super(ResidualBlock, self).__init__()        self.conv1 = nn.Conv2d(channels, channels, kernel_size=3, padding=1)        self.bn1 = nn.BatchNorm2d(channels)        self.prelu = nn.PReLU(channels)        self.conv2 = nn.Conv2d(channels, channels, kernel_size=3, padding=1)        self.bn2 = nn.BatchNorm2d(channels)    def forward(self, x):        short_cut = x        x = self.conv1(x)        x = self.bn1(x)        x = self.prelu(x)        x = self.conv2(x)        x = self.bn2(x)        return x + short_cutclass UpsampleBLock(nn.Module):    def __init__(self, in_channels, up_scale):        super(UpsampleBLock, self).__init__()        self.conv = nn.Conv2d(in_channels, in_channels * up_scale ** 2, kernel_size=3, padding=1)        self.pixel_shuffle = nn.PixelShuffle(up_scale)        self.prelu = nn.PReLU(in_channels)    def forward(self, x):        x = self.conv(x)        x = self.pixel_shuffle(x)        x = self.prelu(x)        return xclass Generator(nn.Module):    def __init__(self, scale_factor, num_residual=16):        upsample_block_num = int(math.log(scale_factor, 2))        super(Generator, self).__init__()        self.block_in = nn.Sequential(            nn.Conv2d(3, 64, kernel_size=9, padding=4),            nn.PReLU(64)        )        self.blocks = []        for _ in range(num_residual):            self.blocks.append(ResidualBlock(64))        self.blocks = nn.Sequential(*self.blocks)        self.block_out = nn.Sequential(            nn.Conv2d(64, 64, kernel_size=3, padding=1),            nn.BatchNorm2d(64)        )        self.upsample = [UpsampleBLock(64, 2) for _ in range(upsample_block_num)]        self.upsample.append(nn.Conv2d(64, 3, kernel_size=9, padding=4))        self.upsample = nn.Sequential(*self.upsample)    def forward(self, x):        x = self.block_in(x)        short_cut = x        x = self.blocks(x)        x = self.block_out(x)        upsample = self.upsample(x + short_cut)        return torch.tanh(upsample)

三、判别网络的构建

Pytorch怎么搭建SRGAN平台提升图片超分辨率

判别网络的构成如上图所示：

SRGAN的判别网络由不断重复的卷积+LeakyRELU和标准化组成。
对于判断网络来讲，它的目的是判断输入图片的真假，它的输入是图片，输出是判断结果。

判断结果处于0-1之间，利用接近1代表判断为真图片，接近0代表判断为假图片。

判断网络的构建和普通卷积网络差距不大，都是不断的卷积对图片进行下采用，在多次卷积后，最终接一次全连接判断结果。

实现代码如下：

class Discriminator(nn.Module):    def __init__(self):        super(Discriminator, self).__init__()        self.net = nn.Sequential(            nn.Conv2d(3, 64, kernel_size=3, padding=1),            nn.LeakyReLU(0.2),            nn.Conv2d(64, 64, kernel_size=3, stride=2, padding=1),            nn.BatchNorm2d(64),            nn.LeakyReLU(0.2),            nn.Conv2d(64, 128, kernel_size=3, padding=1),            nn.BatchNorm2d(128),            nn.LeakyReLU(0.2),            nn.Conv2d(128, 128, kernel_size=3, stride=2, padding=1),            nn.BatchNorm2d(128),            nn.LeakyReLU(0.2),            nn.Conv2d(128, 256, kernel_size=3, padding=1),            nn.BatchNorm2d(256),            nn.LeakyReLU(0.2),            nn.Conv2d(256, 256, kernel_size=3, stride=2, padding=1),            nn.BatchNorm2d(256),            nn.LeakyReLU(0.2),            nn.Conv2d(256, 512, kernel_size=3, padding=1),            nn.BatchNorm2d(512),            nn.LeakyReLU(0.2),            nn.Conv2d(512, 512, kernel_size=3, stride=2, padding=1),            nn.BatchNorm2d(512),            nn.LeakyReLU(0.2),            nn.AdaptiveAvgPool2d(1),            nn.Conv2d(512, 1024, kernel_size=1),            nn.LeakyReLU(0.2),            nn.Conv2d(1024, 1, kernel_size=1)        )    def forward(self, x):        batch_size = x.size(0)        return torch.sigmoid(self.net(x).view(batch_size))