基于通道注意力机制的图像压缩改进（zqb第三章）

小夭。

已于 2023-01-05 22:37:45 修改

阅读量506

点赞数

分类专栏：文献代码复现 TensorFlow 图像压缩文章标签： python 深度学习

于 2022-08-30 14:45:17 首次发布

本文链接：https://blog.csdn.net/m0_47146037/article/details/126603730

版权

本文详细介绍了使用通道注意力机制改进图像压缩的实现过程，包括数据集准备、训练流程、关键模块（如通道注意力模块ATAC、非线性分析变换等）的实现，以及训练和验证中的问题与解决方案。通过对kodak和clic数据集的实验，复现并对比了cheng2020实验数据，展示了不同码率点的性能表现。

摘要由CSDN通过智能技术生成

前言

此篇博客仅供自己观看，记录实验复现中的问题。

1、数据集准备

1.1 拷贝师兄的数据集到自己对应的目录（imageNet database）

在这里插入图片描述
下面为目录的更改

原始拷贝的数据集全是tar压缩包的形式，添加.sh脚本进行批量解压缩tar压缩包。

for i in $(ls *tar);
do tar -xvf $i ;
done

在这里插入图片描述

1.2 直接集群调用师兄的数据集

在这里插入图片描述

2、调用train.py进行训练（通过launch.json暂时无法调试）

同样添加.sh脚本进行训练以及相应日志的输出。
在这里插入图片描述

conda activate zqb
python -u train.py train > log.out

3、整理训练流程

在这里插入图片描述

读取数据集地址，加载数据集
随机剪裁数据中的图片为patchsize*patchsize大小
训练前加载设置好的batchsize，进行图像的分批
batchsize * patchsize² 计算像素数
通过非线性分析变换analysis_transform（g_a）获取隐表示y
通过超先验部分的非线性分析变换hyper_analysis（h_a）获取超隐表示z
调用tfc的熵模型EntropyBottleneck获取经过量化（添加均匀噪声后的值）与对数似然值
通过上一步获取的值进行超先验部分的非线性综合变换hyper_synthesis（h_s）获取自编码器的统计参数φ
设置tiny_y tiny_phi占位符，通过entropy_parameter模块获得均值、方差等统计信息，从而获取y经过经过量化（添加均匀噪声）后的y与对数似然值
通过非线性综合变换synthesis_transform（g_s）重建图像x_hat
通过上述y与z的对数自然概率值与像素值可以计算bpp
通过对比原始图像与重建图像获取mse
进而获取其他性能值如PSNR、MSSSIM
有了bpp与mse，loss = lmbda * mse + bpp
上述为训练一个迭代的过程，将上述过程在优化器中通过设置学习率等参数进行梯度下降从而训练。

4、一些function

4.1、RB残差块（暂时忽略）

在NonLocalAttentionBlock方法中使用，但是后续NonLocalAttentionBlock是在消融实验中用到。
‘’‘后续通过消融实验证实：在相同的压缩框架中，单个 epoch 训练时 RBAM 压缩框架耗时仅为 NLAM 压缩框架的三分之一，且两个框架对损失的贡献几乎相同。综上所述，为了在计算量和性能之间取得更好的平衡，RBAM 相对于 NLAM 更适合作为编解码框架的注意力模块。’‘’

''' 定义一个残差模块residualblock，内含三个2D卷积：1x1 Conv N/2 relu, 3x3 Conv N/2 relu, 1x1 Conv N '''
def residualblock(tensor, num_filters, scope="residual_block"):
  """Builds the residual block"""
  with tf.variable_scope(scope):
    with tf.variable_scope("conv0"):
      layer = tfc.SignalConv2D(
        num_filters//2, (1,1), corr=True, strides_down=1, padding="same_zeros",
        use_bias=True, activation=ChaATAC(channels=num_filters//2), name='signal_conv2d') # 激活函数ChaATAC
      output = layer(tensor)

    with tf.variable_scope("conv1"):
      layer = tfc.SignalConv2D(
        num_filters//2, (3,3), corr=True, strides_down=1, padding="same_zeros",
        use_bias=True, activation=ChaATAC(channels=num_filters//2), name='signal_conv2d')
      output = layer(output)

    with tf.variable_scope("conv2"):
      layer = tfc.SignalConv2D(
        num_filters, (1,1), corr=True, strides_down=1, padding="same_zeros",
        use_bias=True, activation=None, name='signal_conv2d')
      output = layer(output)
      
    tensor = tensor + output
       
  return tensor

在这里插入图片描述

4.2、通道注意力模块 ATAC

注意力模块相关的解释可以见注意力模块
在这里插入图片描述
权重信息

经过权重分配的特征

import tensorflow as tf
from tensorflow.keras.layers import Layer
from tensorflow.keras import activations, layers, Sequential
import tensorflow_compression as tfc

#原始移植代码
# 用法:activation=ChaATAC(channels=num_filters//2), name='signal_conv2d')
class ChaATAC(Layer):
    def __init__(self, channels, r=1, useReLU=True, useGlobal=False):
        super(ChaATAC,self).__init__()
        self.channels = channels
        self.inter_channels = int(channels//r)
        if useGlobal:
            self.avgpool = layers.GlobalAveragePooling2D()
        # 第一层的逐点卷积
        self.conv1 = tfc.SignalConv2D(self.inter_channels, (1,1), corr=True, strides_down=1,
                                    padding='valid',use_bias=False, activation=None, name='Conv1')
        # 第一层批量归一化
        self.bn1 = layers.BatchNormalization()
        # relu激活函数
        if useReLU:
            self.relu = activations.relu
        # 第二层逐点卷积
        self.conv2 = tfc.SignalConv2D(self.channels, (1,1), corr=True, strides_down=1,
                                    padding='valid',use_bias=False, activation=None, name='Conv2')
        # 批量归一化
        self.bn2 = layers.BatchNormalization()
        # sigmoid激活函数
        self.sigmoid = activations.sigmoid

    def call(self, x):
        xs = self.conv1(x)
        xs = self.bn1(xs)
        u = self.relu(xs)

        g = self.conv2(u)
        g = self.bn2(g)
        y = self.sigmoid(g)
        # 最后的y输出对应的权重信息，原特征与权重信息逐点相乘，输出新的特征
        x = x * y
        return x

其实整个代码过程就是按照图片的示意图来进行的。

初始化定义两层的逐点卷积、批归一化、激活函数
在call里分别调用上述定义的方法，获取权重信息y
原始特征x与权重信息y逐点相乘获取新的特征x’

4.3、非线性分析变换analysis_transform

使用三个3x3的卷积，代替一个9x9的卷积。

首先把这个整体分为四个部分，前三个部分分为了四层，第四部分只有一层
在这里插入图片描述
前三个部分的四层其实是一样的

（1）前三部分的每一层

在这里插入图片描述

从代码部分分析，即前三个部分按照四层的网络结构进行编写代码，经过第一二层后会将两层的tensor进行相加，这里就是引入了残差结构，把理想输出逼近输入。剩下两层按照图示结构一步步编写代码。（但是前两层和后两层直接穿插了一个1x1的卷积，不知道为什么）

（2）第四部分

最后一部分就是一个3x3的卷积
在这里插入图片描述

def analysis_transform(tensor, num_filters):
  """Builds the analysis transform."""

  kernel_size = 3
  #Use three 3x3 filters to replace one 9x9
  
  with tf.variable_scope("analysis"):

    # Four down-sampling blocks
    for i in range(4): # 0-3
      # layer0、1、2、3，把四层看成一块，重复三块，加最后一块的3x3卷积
      if i > 0: # 0,1,2,3
        with tf.variable_scope("Block_" + str(i) + "_layer_0"):
          layer = tfc.SignalConv2D(
            num_filters, (kernel_size, kernel_size), corr=True, strides_down=1, padding="same_zeros",
            use_bias=True, activation=ChaATAC(channels=num_filters), name='signal_conv2d')
          tensor2 = layer(tensor)

        with tf.variable_scope("Block_" + str(i) + "_layer_1"):
          layer = tfc.SignalConv2D(
              num_filters, (kernel_size, kernel_size), corr=True, strides_down=1, padding="same_zeros",
              use_bias=True, activation=ChaATAC(channels=num_filters), name='signal_conv2d')
          tensor2 = layer(tensor2)
        
        tensor = tensor + tensor2


      if i < 3: # 0,1,2
        with tf.variable_scope("Block_" + str(i) + "_shortcut"):
          shortcut = tfc.SignalConv2D(num_filters, (1, 1), corr=True, strides_down=2, padding="same_zeros",
                                      use_bias=True, activation=None, name='signal_conv2d')
          shortcut_tensor = shortcut(tensor)

        with tf.variable_scope("Block_" + str(i) + "_layer_2"):
          layer = tfc.SignalConv2D(