基于Paddle复现《Restormer: Efficient Transformer for High-Resolution Image Restoration》

1.简介

由于CNN在从大规模数据中学习广义图像先验知识方面表现良好，这些模型已被广泛应用于图像恢复等相关任务。最近，另一类神经结构Transformers在自然语言和High-Level视觉任务上显示出显著的性能提升。虽然Transformer模型缓解了CNN的不足（即有限的感受野和对输入内容的适应性），但其计算复杂度随空间分辨率呈二次增长，因此无法应用于大多数涉及高分辨率图像的图像恢复任务。在这项工作中，我们提出了一种高效的转换器模型，通过在构建模块（多头注意和前馈网络）中进行几个关键设计，它可以捕获长距离的像素交互，同时仍然适用于大型图像。我们的模型名为RestorationTransformer（Restormer），在多个图像恢复任务上实现了SOTA的结果，本Repo主要复现了图像去噪的模型。

原repo: https://github.com/swz30/Restormer

论文地址: https://arxiv.org/pdf/2111.09881.pdf?ref=https://githubhelp.com

2.复现精度

原repo采用的是8卡训练，这里我改为4卡，同时iters 乘以2，学习率除以2。
在CBSD68测试集的测试效果如下表,达到验收指标,PSNR: 34.39。

Network	opt	iters	learning rate	batch_size	dataset	GPUS	PSNR
Restormer	AdamW	600000	1.5e-4	8	CBSD68	4	34.39

3.数据集

下载地址:

https://aistudio.baidu.com/aistudio/datasetdetail/140244

解压数据集

%cd work/
!cat ../data/data140244/DFWB.tar.gza* | tar zx

最优权重:

链接: https://pan.baidu.com/s/14lxC6gHrr6BXHJBZgY1C_g

提取码: t067

4.环境依赖

PaddlePaddle == 2.2.0

scikit-image == 0.19.2

5.代码解读

数据集部分

数据集为DIV2K, Flickr2K, WED, BSD这4个数据集的融合，简称为DFWB。只需要读取真值图片，样本图片基于真值图片叠加随机噪声生成。

if self.sigma_type == 'constant':
    sigma_value = self.sigma_range
elif self.sigma_type == 'random':
    sigma_value = random.uniform(self.sigma_range[0], self.sigma_range[1])
elif self.sigma_type == 'choice':
    sigma_value = random.choice(self.sigma_range)

noise_level = sigma_value / 255.0
noise = paddle.randn(img_lq.shape,dtype='float32').numpy() * noise_level

img_lq = img_lq + noise.astype('float32')

上述代码根据sigma_type来判断使用哪种方式来获取sigma_value。本项目训练时使用的random，同时sigma_range的范围是0到50。测试的时候使用固定的sigma_test值生成测试样本。

直接使用paddle的randn生成随机噪声，然后与噪声等级相乘，最后叠加到原图上，供模型训练与测试。

模型部分

模型采用Transform的形式构成，下面简单介绍一下。

模型架构如上图所示，其中Transformer模块代码如下：

class TransformerBlock(nn.Layer):
    def __init__(self, dim, num_heads, ffn_expansion_factor, bias, LayerNorm_type):
        super(TransformerBlock, self).__init__()

        self.norm1 = LayerNorm(dim, LayerNorm_type)
        self.attn = Attention(dim, num_heads, bias)
        self.norm2 = LayerNorm(dim, LayerNorm_type)
        self.ffn = FeedForward(dim, ffn_expansion_factor, bias)

    def forward(self, x):
        x = x + self.attn(self.norm1(x))
        x = x + self.ffn(self.norm2(x))

        return x

每个TransformBlock模块还包含了一个Attention和FeedForward模块。代码分别如下：

class Attention(nn.Layer):
    def __init__(self, dim, num_heads, bias):
        super(Attention, self).__init__()
        self.num_heads = num_heads
        # self.temperature = nn.Parameter(torch.ones(num_heads, 1, 1))
        self.temperature = paddle.create_parameter(shape=[num_heads, 1, 1],dtype='float32',
                                              default_initializer=nn.initializer.Constant(1.0))

        self.qkv = nn.Conv2D(dim, dim*3, kernel_size=1, bias_attr=bias)
        self.qkv_dwconv = nn.Conv2D(dim*3, dim*3, kernel_size=3, stride=1, padding=1, groups=dim*3, bias_attr=bias)
        self.project_out = nn.Conv2D(dim, dim, kernel_size=1, bias_attr=bias)
        


    def forward(self, x):
        b,c,h,w = x.shape

        qkv = self.qkv_dwconv(self.qkv(x))
        q,k,v = qkv.chunk(3, axis=1)

        b1, hc, h1, w1 = q.shape
        # q = paddle.reshape(q, [b1, self.num_heads, -1, h1, w1])
        c = hc // self.num_heads
        q = paddle.reshape(q, [b1, self.num_heads, c, (h1*w1)])

        # q = rearrange(q, 'b (head c) h w -> b head c (h w)', head=self.num_heads)
        b1, hc, h1, w1 = k.shape
        # k = paddle.reshape(k, [b1, self.num_heads, -1, h1, w1])
        c = hc // self.num_heads
        k = paddle.reshape(k, [b1, self.num_heads, c, (h1*w1)])
        # k = rearrange(k, 'b (head c) h w -> b head c (h w)', head=self.num_heads)

        b1, hc, h1, w1 = v.shape
        # v = paddle.reshape(v, [b1, self.num_heads, -1, h1, w1])
        c = hc // self.num_heads
        v = paddle.reshape(v, [b1, self.num_heads, c, (h1*w1)])
        # v = rearrange(v, 'b (head c) h w -> b head c (h w)', head=self.num_heads)

        q = paddle.nn.functional.normalize(q, axis=-1)
        k = paddle.nn.functional.normalize(k, axis=-1)

        attn = (q @ k.transpose([0, 1, 3, 2])) * self.temperature
        attn = F.softmax(attn, axis=-1)

        out = (attn @ v)

        b, head, c, hw = out.shape

        # out = rearrange(out, 'b head c (h w) -> b (head c) h w', head=self.num_heads, h=h, w=w)
        # out = paddle.reshape(out, [b, head, c, h, w])
        out = paddle.reshape(out, [b, head * c, h, w])

        out = self.project_out(out)
        return out

其中qkv矩阵由qkv_dwconv卷积后分割得到。qkv_dwconv卷积输出的通道数是输入的通道数的三倍,使用chunk方法后分割为Q、K、V三个三个矩阵。然后计算score=softmax(Q@K)，然后在乘以V值得到输出结果。

以下是FeedForward代码:

class FeedForward(nn.Layer):
    def __init__(self, dim, ffn_expansion_factor, bias):
        super(FeedForward, self).__init__()

        hidden_features = int(dim*ffn_expansion_factor)

        self.project_in = nn.Conv2D(dim, hidden_features*2, kernel_size=1, bias_attr=bias)

        self.dwconv = nn.Conv2D(hidden_features*2, hidden_features*2, kernel_size=3, stride=1, padding=1, groups=hidden_features*2, bias_attr=bias)

        self.project_out = nn.Conv2D(hidden_features, dim, kernel_size=1, bias_attr=bias)

    def forward(self, x):
        x = self.project_in(x)
        x1, x2 = self.dwconv(x).chunk(2, axis=1)
        x = F.gelu(x1) * x2
        x = self.project_out(x)
        return x

该部分代码可以执行可控特征变换，即抑制低信息特征，仅保留有用信息。首先通过dwconv卷积分成两个分支，对其中一个分支进行gelu运算，然后与另外一个分支逐像素相乘实现gating机制，该机制类似空间和通道注意力结合，使网络可以针对每个channel和每个空间位置，学习一种动态的特征选择机制。

6.快速开始

模型训练

训练至少需要4卡资源，配置默认为4卡，如需8卡训练可修改configs/GaussianColorDenoising_Restormer.yml文件。将其中跟iters相关的数值除以2，同时将学习率相关数值乘以2.
多卡训练，启动方式如下：

多卡训练，启动方式如下：

python -u -m paddle.distributed.launch  train.py -opt configs/GaussianColorDenoising_Restormer.yml 

多卡恢复训练，启动方式如下：

python -u -m paddle.distributed.launch  train.py -opt configs/GaussianColorDenoising_Restormer.yml --resume ../245_model

本项目使用脚本模式进行训练，若希望使用notebook模式，请选择4卡v100的环境进行训练，首先安装依赖包，再执行训练命令：

!pip install scikit-image

开始训练，

%cd /home/aistudio/Restormer_Paddle/
!python -u -m paddle.distributed.launch  train.py -opt configs/GaussianColorDenoising_Restormer.yml 

参数介绍：

opt: 配置路径

resume: 从哪个模型开始恢复训练，需要pdparams和pdopt文件。

模型验证

除了可以再训练过程中验证模型精度，还可以是val.py脚本加载模型验证精度，执行以下命令。
验证数据的地址需要设置configs/GaussianColorDenoising_Restormer.yml中的datasets.val.dataroot_gt参数。

%cd /home/aistudio/Restormer_Paddle/
!python val.py -opt configs/GaussianColorDenoising_Restormer.yml --weights best_model.pdparams --sigmas 15 

参数说明：

opt: 配置路径

weights: 模型权重地址

sigmas: 噪声等级

单张图片预测

本项目提供了单张图片的预测脚本，可根据输入图片生成噪声，然后对图片进行降噪。会在result_dir指定的目录下生成denoise_0000.png和noise_0000.png两张图片。使用方法如下：

!pip install natsort
%cd /home/aistudio/Restormer_Paddle/
!python predict.py --input_images demo/0000.png \
--weights best_model.pdparams \
--model_type blind --sigmas 15 --result_dir ./output/

参数说明：

input_images:需要预测的图片

weights: 模型路径

result_dir: 输出图片保存路径

model_type: 模型类型，本项目只训练了blind模式。

sigmas: 噪声等级。

在噪声等级15下的预测样例:

模型导出

模型导出可执行以下命令：

%cd /home/aistudio/Restormer_Paddle/
!python export_model.py -opt ./configs/GaussianColorDenoising_Restormer.yml --model_path ./output/model/last_model.pdparams --save_dir ./output/

参数说明：

opt: 模型配置路径

model_path: 模型路径

save_dir: 输出图片保存路径

Inference推理

可使用以下命令进行模型推理。该脚本依赖auto_log, 请参考下面TIPC部分先安装auto_log。infer命令运行如下：

%cd /home/aistudio/Restormer_Paddle/
!python infer.py --use_gpu=False --enable_mkldnn=False --cpu_threads=2 --model_file=./test_tipc/output/model.pdmodel --batch_size=2 --input_file=test_tipc/data/CBSD68 --enable_benchmark=True --precision=fp32 --params_file=./test_tipc/output/model.pdiparams 

参数说明:

use_gpu:是否使用GPU

enable_mkldnn:是否使用mkldnn

cpu_threads: cpu线程数

model_file: 模型路径

batch_size: 批次大小

input_file: 输入文件路径

enable_benchmark: 是否开启benchmark

precision: 运算精度

params_file: 模型权重文件，由export_model.py脚本导出。

TIPC基础链条测试

该部分依赖auto_log，需要进行安装，安装方式如下：

auto_log的详细介绍参考https://github.com/LDOUBLEV/AutoLog。

%cd /home/aistudio/
!git clone https://gitee.com/Double_V/AutoLog
%cd AutoLog/
!pip install -r requirements.txt
!python setup.py bdist_wheel
!pip install ./dist/auto_log-1.2.0-py3-none-any.whl

%cd /home/aistudio/Restormer_Paddle/
!bash test_tipc/prepare.sh ./test_tipc/configs/Restormer/train_infer_python.txt 'lite_train_lite_infer'

!bash test_tipc/test_train_inference_python.sh ./test_tipc/configs/Restormer/train_infer_python.txt 'lite_train_lite_infer'

测试结果如截图所示：

7.代码结构与详细说明

Restormer_Paddle
├── README.md # 说明文件
├── logs # 训练日志
├── configs # 配置文件
├── data # 数据变换
├── dataset.py # 数据集路径
├── demo # 样例图片
├── export_model.py # 模型导出
├── infer.py  # 推理预测
├── metrics  # 指标计算方法
├── models # 网络模型
├── predict.py # 图像预测
├── test_tipc # TIPC测试链条
├── train.py # 训练脚本
├── utils # 工具类
└── val.py # 评估脚本

8.模型信息

信息	描述
模型名称	Restormer
框架版本	PaddlePaddle==2.2.0
应用场景	降噪

9.心得体会

复现思路

首先阅读论文和原repo，重点读懂原repo的代码，通过Readme可了解到，达到验收指标，使用的是GaussianColorDenoising_Restormer.yml配置文件。根据配置文件内容找到相关的数据集和模型网络代码并复现。其中原项目网络使用torch实现，大部分接口paddle都有对应的实现，直接替换即可，其中需要注意的是，网络中使用了以下代码对张量进行转置。

q = rearrange(q, 'b (head c) h w -> b head c (h w)', head=self.num_heads)

rearrange并没有paddle实现，所以针对类似的代码需要使用reshape来替代，同时还需要兼顾模型导出时，reshape中的shape参数不能包含两个-1，所以修改完的代码如下：

b1, hc, h1, w1 = q.shape
c = hc // self.num_heads
q = paddle.reshape(q, [b1, self.num_heads, c, (h1*w1)])

网络复现完成后，就与其他的训练方式基本一致，构建dataloader、net以及optimizer，编写循环代码，设置循环终止条件，并编写日志输出策略、模型验证及保存策略。这里需要注意的是，由于论文中要求的batch size比较大，原repo使用的是8卡进行训练。为了能快速复现我们也需要编写多卡训练的代码，可参考PaddleSeg等套件中的实现，添加了多卡训练，同时注意保存模型和日志输出只在0号卡上实现。由于AI Studio最多只有4卡资源，所以相当于batch size是原repo的一半，这样训练的迭代次数就需要增加一倍，同时learning rate需要改成原来的一半，训练时间延长一倍。

心得体会

由于数据集和模型规模比较大，本次复现是我单次训练时间最长的一次，原Repo使用的是8卡资源，我使用的是4卡V100，所以训练时间比原Repo应该是多了一倍，一共花费了12天的时间，其中包括脚本任务排队的时间。以前很少尝试这种大模型和大数据集的复现，通过这次项目掌握了编写多卡训练程序的方法，以后遇到这种大规模的模型复现也有了信心，同时也要感谢AI Studio提供的算力支持，如果没有算力，自己是无法支持这么大规模训练的。