5、盲超分辨-KernelGAN

安逸sgr

于 2024-10-04 10:36:47 发布

阅读量882

点赞数 27

分类专栏：对抗生成网络图像生成与修复文章标签：计算机视觉深度学习神经网络生成对抗网络 pytorch 图像处理超分辨率重建

本文链接：https://blog.csdn.net/sgr011215/article/details/142700635

版权

对抗生成网络同时被 2 个专栏收录

6 篇文章 0 订阅

订阅专栏

图像生成与修复

4 篇文章 0 订阅

订阅专栏

5、KernelGAN

KernelGAN论文链接：Blind Super-Resolution Kernel Estimation using an Internal-GAN

KernelGAN的主要创新点在于它引入了一个 图像特定的内部生成对抗网络 （Internal GAN），用于解决盲超分辨率（Blind Super-Resolution）中的卷积核估计问题。

模型整体架构图

主要思想：

首先在原始图像中裁剪一个区域作为真实补丁，后续用来和的生成补丁计算Loss值来更新网络
生成器，对于生成器使用卷鸡和线性变换组成，没有非线性变换，防止图像发生本质性的变化，生成出缩小后的图像，同时裁剪出和真实补丁一样大小的一块送入到判别器。
判别器 ，采用一个类似于热力图的概率分布图D-map来实现计算每个像素点之间的差距，也就是均方误差 (MSE) ，将真实补丁判别为真，将生成的补丁判别为假，其中D-map每个值就每个像素点为真的概率，对于真实补丁，D-map 的输出应接近 1；对于生成补丁，D-map 的输出应接近 0。
损失函数，不仅是判别器的损失函数，我们还需要生成一个图像特定的卷积核，这是最终的目标，所以还需要一些控制生成标准的卷积核损失函数。
提取显式内核，主要是通过对生成器的卷积核的参数进行卷积生成最终的显式内核。

在这里插入图片描述

生成器

生成器架构图，生成器 G 被设计为一种线性下采样模型，它通过卷积和下采样对图像进行线性变换，用来逼近正确的下采样核。作者认为单层卷积生成器理论上能够表达所有可能的下采样方法，但在实践中它无法收敛到正确的解决方案。

生成器 G 包含 5 个线性卷积层，每一层有 64 个通道，滤波器的大小分别是 7×7、5×5、3×3，最后两层是 1×1 的卷积核。这种设计允许生成器覆盖 13×13 的感受野，从而生成一个 13×13 的下采样卷积核。

在这里插入图片描述

class Generator(nn.Module):
    def __init__(self, conf):
        super(Generator, self).__init__()
        struct = conf.G_structure
        
        # 第一层 - 将 RGB 图像转换为潜在空间
        self.first_layer = nn.Conv2d(in_channels=1, out_channels=conf.G_chan, kernel_size=struct[0], bias=False)

        # 特征提取模块 - 中间层的堆叠
        feature_block = []  # 使用多个卷积层
        for layer in range(1, len(struct) - 1):
            feature_block += [
                nn.Conv2d(in_channels=conf.G_chan, out_channels=conf.G_chan, kernel_size=struct[layer], bias=False)
            ]
        self.feature_block = nn.Sequential(*feature_block)
        
        # 最后一层 - 降采样并转换回图像
        self.final_layer = nn.Conv2d(
            in_channels=conf.G_chan, out_channels=1, kernel_size=struct[-1],
            stride=int(1 / conf.scale_factor), bias=False
        )

        # 计算在前向传播时削减的像素数
        self.output_size = self.forward(torch.FloatTensor(torch.ones([1, 1, conf.input_crop_size, conf.input_crop_size]))).shape[-1]
        self.forward_shave = int(conf.input_crop_size * conf.scale_factor) - self.output_size

    def forward(self, input_tensor):
        # 调整 RGB 图像的轴，使网络能够处理大小为 3 的 "batch" 而不是 3 个通道
        input_tensor = swap_axis(input_tensor)  # 假设 swap_axis 是一个函数，用于调整张量的轴
        
        # 首先将输入图像传递给第一层
        downscaled = self.first_layer(input_tensor)
        
        # 通过特征提取模块获取特征
        features = self.feature_block(downscaled)
        
        # 最后一层将特征转换回图像
        output = self.final_layer(features)
        
        # 调整输出图像的轴
        return swap_axis(output)

判别器

判别器 D 学习输入图像的局部补丁分布，并通过输出热力图来区分图像中的真实和虚假补丁。它采用完全卷积结构，使用 7×7 的小感受野，没有池化和步幅，确保 D 能够对每个局部补丁进行单独判断，并输出每个位置的概率。这种设计有助于判别器细致地区分生成的虚假图像补丁与真实图像补丁的差异。

假设有一个 32×32 的低分辨率图像输入到判别器 D。该图像被分为多个 7×7 的小块补丁，每个补丁都是 D 需要判断的对象。判别器的任务是区分这些 7×7 补丁来自真实图像 ILRI_{LR}ILR 还是生成器 G 生成的图像。

如果判别器中有一个卷积层使用 7×7 的卷积核（感受野），它会对每个 7×7 补丁应用卷积操作，输出的每个位置就会对应一个补丁的真假判断。

在这里插入图片描述

class Discriminator(nn.Module):

    def __init__(self, conf):
        super(Discriminator, self).__init__()

        # 第一层 - 卷积层 (没有使用 ReLU 激活函数)
        # 使用光谱归一化 (spectral normalization) 来增强模型的稳定性
        self.first_layer = nn.utils.spectral_norm(
            nn.Conv2d(in_channels=3, out_channels=conf.D_chan, kernel_size=conf.D_kernel_size, bias=True)
        )
        
        # 特征提取模块 (由多个 1x1 卷积层组成)
        feature_block = []  # 使用 1x1 卷积层堆叠层数
        for _ in range(1, conf.D_n_layers - 1):
            feature_block += [
                nn.utils.spectral_norm(nn.Conv2d(in_channels=conf.D_chan, out_channels=conf.D_chan, kernel_size=1, bias=True)),
                nn.BatchNorm2d(conf.D_chan),  # 批归一化层，确保输出分布的稳定
                nn.ReLU(True)  # 激活函数
            ]
        self.feature_block = nn.Sequential(*feature_block)
        
        # 最后一层 - 1x1 卷积层和 Sigmoid 激活函数
        # 输出 1 表示每个图像 patch 的真假
        self.final_layer = nn.Sequential(
            nn.utils.spectral_norm(nn.Conv2d(in_channels=conf.D_chan, out_channels=1, kernel_size=1, bias=True)),
            nn.Sigmoid()  # Sigmoid 函数将输出映射到 (0,1) 区间
        )

        # 计算在前向传递时削减的像素数
        # 通过模拟一次正向传播来确定裁剪大小
        self.forward_shave = conf.input_crop_size - self.forward(
            torch.FloatTensor(torch.ones([1, 3, conf.input_crop_size, conf.input_crop_size]))
        ).shape[-1]

    def forward(self, input_tensor):
        # 提取感受野的特征（第一层卷积）
        receptive_extraction = self.first_layer(input_tensor)
        
        # 通过特征提取模块 (1x1 卷积 + 批归一化 + ReLU)
        features = self.feature_block(receptive_extraction)
        
        # 最后一层生成真假判别的概率图
        return self.final_layer(features)

损失函数

主要介绍生成器损失函数，因为特定的卷积核就是通过生成器的卷积核参数来提取的

这个损失函数表示，生成器 G和判别器 D 之间的目标是互相对抗的，生成器希望生成的图像可以欺骗判别器，而判别器则尝试区分生成图像和真实图像。最终，生成器的优化目的是通过调整生成的图像，使其尽可能接近真实图像。

主损失函数

loss_g 对应前面对抗损失，self.calc_constraints(g_pred)对应后面的R损失，也就是控制生成标准核的损失函数。

在这里插入图片描述

def train_g(self):
    # 清空梯度
    self.optimizer_G.zero_grad()
    
    # 生成器前向传播，生成假图像
    g_pred = self.G.forward(self.g_input)
    
    # 将生成器的输出传递给判别器进行判断
    d_pred_fake = self.D.forward(g_pred)
    
    # 计算生成器损失，基于判别器对生成器输出的预测
    # is_d_input_real=True 表示我们在计算生成器损失时，认为输入的图像是“真实”的
    loss_g = self.criterionGAN(d_last_layer=d_pred_fake, is_d_input_real=True)
    
    # 计算所有损失的总和，包括生成器损失和其他约束损失 loss_g
    total_loss_g = loss_g + self.calc_constraints(g_pred)
    
    # 计算梯度
    total_loss_g.backward()
    
    # 更新生成器的权重
    self.optimizer_G.step()

控制生成标准核的损失函数

通过四个损失来控制生成符合标准的卷积核。

在这里插入图片描述

这些项组合在一起，帮助模型学习到一个有中心化、稀疏、并且归一化的卷积核，用于图像处理或生成任务中。

def calc_constraints(self, g_pred):
    # 计算与生成器 G 相当的卷积核 K
    self.calc_curr_k()
    
    # 计算约束损失
    # 计算双线性插值损失
    self.loss_bicubic = self.bicubic_loss.forward(g_input=self.g_input, g_output=g_pred)
    
    # 计算边界损失
    loss_boundaries = self.boundaries_loss.forward(kernel=self.curr_k)
    
    # 计算归一化损失，使卷积核的和为 1
    loss_sum2one = self.sum2one_loss.forward(kernel=self.curr_k)
    
    # 计算中心化损失，确保卷积核的重心位于中心
    loss_centralized = self.centralized_loss.forward(kernel=self.curr_k)
    
    # 计算稀疏损失，鼓励卷积核稀疏以防止过度平滑
    loss_sparse = self.sparse_loss.forward(kernel=self.curr_k)
    
    # 应用约束系数，计算所有约束损失的加权和
    return self.loss_bicubic * self.lambda_bicubic + \
           loss_sum2one * self.lambda_sum2one + \
           loss_boundaries * self.lambda_boundaries + \
           loss_centralized * self.lambda_centralized + \
           loss_sparse * self.lambda_spar

提取显式内核

主要是将生成器G的每一层参数进行卷积

对于第一层参数，我们初始化一张全为1的特征图delta，让生成器第一层卷积核对delta进行卷积，实际上就是将原始的特征图进行复制得到delta1。
后面的层的卷积和对前面卷积得到的特征图delta1进行卷积，最终得到最后的特征图，也就是我们最后需要的显式内核k。

def calc_curr_k(self):
    """给定生成器网络，该函数计算它正在模仿的卷积核"""
    # 创建一个张量 delta，形状为 (1, 1, 1, 1)，用于卷积操作
    delta = torch.Tensor([1.]).unsqueeze(0).unsqueeze(-1).unsqueeze(-1).cuda()
    
    # 遍历生成器 G 的所有参数（卷积核）
    for ind, w in enumerate(self.G.parameters()):
        # 如果是第一个卷积层，则将 delta 与当前卷积核 w 进行卷积
        # 否则，将前一步的卷积结果 curr_k 与当前卷积核 w 进行卷积
        curr_k = F.conv2d(delta, w, padding=self.conf.G_kernel_size - 1) if ind == 0 else F.conv2d(curr_k, w)
    
    # 将 curr_k 压缩维度，并进行翻转，使得卷积核的顺序符合要求
    self.curr_k = curr_k.squeeze().flip([0, 1])

有个小问题

计算相同大小的补丁时，缩小后的图像截取相同大小的补丁，对于缩小后的图像截取的区域会更大的，那跟真实补丁比较，这是为什么呢？

在 KernelGAN 中，即使缩小后的图像整体变小，但我们仍然从中提取与原始低分辨率图像相同大小的局部补丁，这看似存在比较不一致的潜在问题。那么，为什么可以进行这样的对比呢？这里涉及到两个重要概念：

跨尺度自相似性

KernelGAN 依赖于一个关键的图像属性：跨尺度的局部补丁自相似性。这种自相似性是指，在自然图像中，小范围的局部图像特征（例如纹理、边缘等）在不同尺度上是重复出现的。即便图像经过缩小或放大，局部补丁的形状、纹理特征仍会保留相似性。

因此，即使生成器输出的是一个经过缩小的图像，在这个图像中的局部补丁仍然可以保持与原始图像补丁相似的特征。这为补丁之间的对比提供了合理的基础。

感受野的扩展

感受野（Receptive Field）是指网络中的某个神经元能够“看见”的输入区域。在卷积神经网络中，感受野随着网络层数的增加会逐渐扩展。因此，虽然图像缩小了，但是网络能够提取更大范围的特征。

当我们从缩小的图像中提取相同大小的补丁时，实际上感受野会更大，涵盖到更大范围的图像信息。因此，虽然缩小后的图像的实际像素数量减少，但从网络提取的补丁包含了相应范围的特征，这使得缩小后的图像的局部补丁仍可以用于与原始低分辨率图像的补丁进行对比。

Loss 计算的意义

在这样的对比过程中，判别器在比较的是不同尺度下的局部特征分布是否一致，而不是图像的全局大小。因此，即使缩小后的图像提取的区域相对较大，判别器仍然能够学习到这些特征分布之间的相似性。

最终，生成器通过不断调整其输出图像的下采样核，使得在不同尺度下的图像局部特征保持一致，从而达到优化的效果。