基于Paddle实现Transformer图像降噪

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

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

论文地址: https://arxiv.org/pdf/2111.09881.pdf?ref&＃61;https://githubhelp.com

2.复现精度
原repo采用的是8卡训练&＃xff0c;这里我改为4卡&＃xff0c;同时iters 乘以2&＃xff0c;学习率除以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

解压数据集

In [ ]
%cd data/
!cat data140244/DFWB.tar.gza* | tar zxv
最优权重:

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

提取码: t067

4.环境依赖
PaddlePaddle &＃61;&＃61; 2.2.0

scikit-image &＃61;&＃61; 0.19.2

5.代码解读
数据集部分
数据集为DIV2K, Flickr2K, WED, BSD这4个数据集的融合&＃xff0c;简称为DFWB。只需要读取真值图片&＃xff0c;样本图片基于真值图片叠加随机噪声生成。

if self.sigma_type &＃61;&＃61; ‘constant’:
sigma_value &＃61; self.sigma_range
elif self.sigma_type &＃61;&＃61; ‘random’:
sigma_value &＃61; random.uniform(self.sigma_range[0], self.sigma_range[1])
elif self.sigma_type &＃61;&＃61; ‘choice’:
sigma_value &＃61; random.choice(self.sigma_range)

noise_level &＃61; sigma_value / 255.0
noise &＃61; paddle.randn(img_lq.shape,dtype&＃61;‘float32’).numpy() * noise_level

img_lq &＃61; img_lq &＃43; noise.astype(‘float32’)
上述代码根据sigma_type来判断使用哪种方式来获取sigma_value。本项目训练时使用的random&＃xff0c;同时sigma_range的范围是0到50。测试的时候使用固定的sigma_test值生成测试样本。

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

模型部分
模型采用Transform的形式构成&＃xff0c;下面简单介绍一下。
模型架构如上图所示&＃xff0c;其中Transformer模块代码如下&＃xff1a;

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

self.norm1 &＃61; LayerNorm(dim, LayerNorm_type)self.attn &＃61; Attention(dim, num_heads, bias)self.norm2 &＃61; LayerNorm(dim, LayerNorm_type)self.ffn &＃61; FeedForward(dim, ffn_expansion_factor, bias)def forward(self, x):x &＃61; x &＃43; self.attn(self.norm1(x))x &＃61; x &＃43; self.ffn(self.norm2(x))return x


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

class Attention(nn.Layer):
def init(self, dim, num_heads, bias):
super(Attention, self).init()
self.num_heads &＃61; num_heads
# self.temperature &＃61; nn.Parameter(torch.ones(num_heads, 1, 1))
self.temperature &＃61; paddle.create_parameter(shape&＃61;[num_heads, 1, 1],dtype&＃61;‘float32’,
default_initializer&＃61;nn.initializer.Constant(1.0))

self.qkv &＃61; nn.Conv2D(dim, dim*3, kernel_size&＃61;1, bias_attr&＃61;bias)self.qkv_dwconv &＃61; nn.Conv2D(dim*3, dim*3, kernel_size&＃61;3, stride&＃61;1, padding&＃61;1, groups&＃61;dim*3, bias_attr&＃61;bias)self.project_out &＃61; nn.Conv2D(dim, dim, kernel_size&＃61;1, bias_attr&＃61;bias)def forward(self, x):b,c,h,w &＃61; x.shapeqkv &＃61; self.qkv_dwconv(self.qkv(x))q,k,v &＃61; qkv.chunk(3, axis&＃61;1)b1, hc, h1, w1 &＃61; q.shape# q &＃61; paddle.reshape(q, [b1, self.num_heads, -1, h1, w1])c &＃61; hc // self.num_headsq &＃61; paddle.reshape(q, [b1, self.num_heads, c, (h1*w1)])# q &＃61; rearrange(q, &＃39;b (head c) h w -> b head c (h w)&＃39;, head&＃61;self.num_heads)b1, hc, h1, w1 &＃61; k.shape# k &＃61; paddle.reshape(k, [b1, self.num_heads, -1, h1, w1])c &＃61; hc // self.num_headsk &＃61; paddle.reshape(k, [b1, self.num_heads, c, (h1*w1)])# k &＃61; rearrange(k, &＃39;b (head c) h w -> b head c (h w)&＃39;, head&＃61;self.num_heads)b1, hc, h1, w1 &＃61; v.shape# v &＃61; paddle.reshape(v, [b1, self.num_heads, -1, h1, w1])c &＃61; hc // self.num_headsv &＃61; paddle.reshape(v, [b1, self.num_heads, c, (h1*w1)])# v &＃61; rearrange(v, &＃39;b (head c) h w -> b head c (h w)&＃39;, head&＃61;self.num_heads)q &＃61; paddle.nn.functional.normalize(q, axis&＃61;-1)k &＃61; paddle.nn.functional.normalize(k, axis&＃61;-1)attn &＃61; (q &＃64; k.transpose([0, 1, 3, 2])) * self.temperatureattn &＃61; F.softmax(attn, axis&＃61;-1)out &＃61; (attn &＃64; v)b, head, c, hw &＃61; out.shape# out &＃61; rearrange(out, &＃39;b head c (h w) -> b (head c) h w&＃39;, head&＃61;self.num_heads, h&＃61;h, w&＃61;w)# out &＃61; paddle.reshape(out, [b, head, c, h, w])out &＃61; paddle.reshape(out, [b, head * c, h, w])out &＃61; self.project_out(out)return out


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

以下是FeedForward代码:

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

hidden_features &＃61; int(dim*ffn_expansion_factor)self.project_in &＃61; nn.Conv2D(dim, hidden_features*2, kernel_size&＃61;1, bias_attr&＃61;bias)self.dwconv &＃61; nn.Conv2D(hidden_features*2, hidden_features*2, kernel_size&＃61;3, stride&＃61;1, padding&＃61;1, groups&＃61;hidden_features*2, bias_attr&＃61;bias)self.project_out &＃61; nn.Conv2D(hidden_features, dim, kernel_size&＃61;1, bias_attr&＃61;bias)def forward(self, x):x &＃61; self.project_in(x)x1, x2 &＃61; self.dwconv(x).chunk(2, axis&＃61;1)x &＃61; F.gelu(x1) * x2x &＃61; self.project_out(x)return x


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

6.快速开始
模型训练
训练至少需要4卡资源&＃xff0c;配置默认为4卡&＃xff0c;如需8卡训练可修改configs/GaussianColorDenoising_Restormer.yml文件。将其中跟iters相关的数值除以2&＃xff0c;同时将学习率相关数值乘以2. 多卡训练&＃xff0c;启动方式如下&＃xff1a;

多卡训练&＃xff0c;启动方式如下&＃xff1a;

python -u -m paddle.distributed.launch train.py -opt configs/GaussianColorDenoising_Restormer.yml
多卡恢复训练&＃xff0c;启动方式如下&＃xff1a;

python -u -m paddle.distributed.launch train.py -opt configs/GaussianColorDenoising_Restormer.yml --resume …/245_model
本项目使用脚本模式进行训练&＃xff0c;若希望使用notebook模式&＃xff0c;请选择4卡v100的环境进行训练&＃xff0c;可执行以下命令&＃xff1a;

In [ ]
!pip install scikit-image
%cd /home/aistudio/Restormer_Paddle/
!python -u -m paddle.distributed.launch train.py -opt configs/GaussianColorDenoising_Restormer.yml
参数介绍&＃xff1a;

opt: 配置路径

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

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

In [ ]
%cd /home/aistudio/Restormer_Paddle/
!python val.py -opt configs/GaussianColorDenoising_Restormer.yml --weights best_model.pdparams --sigmas 15
参数说明&＃xff1a;

opt: 配置路径

weights: 模型权重地址

sigmas: 噪声等级

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

In [7]
!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/
Looking in indexes: https://pypi.tuna.tsinghua.edu.cn/simple
Collecting natsort
Downloading https://pypi.tuna.tsinghua.edu.cn/packages/a9/76/0f624b7326f4458a249580c55e5654756084ec4572ce37a05f799b96bc24/natsort-8.1.0-py3-none-any.whl (37 kB)
Installing collected packages: natsort
Successfully installed natsort-8.1.0
WARNING: You are using pip version 22.0.4; however, version 22.1.2 is available.
You should consider upgrading via the ‘/opt/conda/envs/python35-paddle120-env/bin/python -m pip install --upgrade pip’ command.
/home/aistudio/Restormer_Paddle
/opt/conda/envs/python35-paddle120-env/lib/python3.7/site-packages/setuptools/depends.py:2: DeprecationWarning: the imp module is deprecated in favour of importlib; see the module’s documentation for alternative uses
import imp
Compute results for noise level 15
Loading pretrained model from best_model.pdparams
There are 406/406 variables loaded into Restormer.
&＃61;&＃61;&＃61;>Testing using weights: best_model.pdparams

100%|█████████████████████████████████████████████| 1/1 [00:44<00:00, 44.79s/it]
The predict image save in ./output/ path.
参数说明&＃xff1a;

input_images:需要预测的图片

weights: 模型路径

result_dir: 输出图片保存路径

model_type: 模型类型&＃xff0c;本项目只训练了blind模式。

sigmas: 噪声等级。

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

从左到右分别是clear、nosie、denoise
模型导出
模型导出可执行以下命令&＃xff1a;

In [ ]
%cd /home/aistudio/Restormer_Paddle/
!python export_model.py -opt ./test_tipc/configs/GaussianColorDenoising_Restormer.yml --model_path ./output/model/last_model.pdparams --save_dir ./test_tipc/output/
参数说明&＃xff1a;

opt: 模型配置路径

model_path: 模型路径

save_dir: 输出图片保存路径

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

In [ ]
%cd /home/aistudio/Restormer_Paddle/
!python infer.py --use_gpu&＃61;False --enable_mkldnn&＃61;False --cpu_threads&＃61;2 --model_file&＃61;./test_tipc/output/model.pdmodel --batch_size&＃61;2 --input_file&＃61;test_tipc/data/CBSD68 --enable_benchmark&＃61;True --precision&＃61;fp32 --params_file&＃61;./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: 模型权重文件&＃xff0c;由export_model.py脚本导出。

TIPC基础链条测试
该部分依赖auto_log&＃xff0c;需要进行安装&＃xff0c;安装方式如下&＃xff1a;

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

git clone https://gitee.com/Double_V/AutoLog
cd AutoLog/
pip3 install -r requirements.txt
python3 setup.py bdist_wheel
pip3 install ./dist/auto_log-1.2.0-py3-none-any.whl
In [ ]
%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’
测试结果如截图所示&＃xff1a;
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&＃61;&＃61;2.2.0
应用场景 降噪
9.心得体会
由于数据集和模型规模比较大&＃xff0c;本次复现是我单次训练时间最长的一次&＃xff0c;原Repo使用的是8卡资源&＃xff0c;我使用的是4卡V100&＃xff0c;所以训练时间比原Repo应该是多了一倍&＃xff0c;一共花费了12天的时间&＃xff0c;其中包括脚本任务排队的时间。以前很少尝试这种大模型和大数据集的复现&＃xff0c;通过这次项目掌握了编写多卡训练程序的方法&＃xff0c;以后遇到这种大规模的模型复现也有了信心&＃xff0c;同时也要感谢AI Studio提供的算力支持&＃xff0c;如果没有算力&＃xff0c;自己是无法支持这么大规模训练的。