421zuoduan / Project-for-Statistical-Learning

Derain Methods Based on Uformer

Project for Statistical Learning Course

Ruo-Chen Cui

Abstract: Our project is based on Restormer, (write later)

---

Related Work

Restormer : Paper | Code

SENet : Paper | Code

Uformer : Paper | Code

Our training framework : Code

Our Work

我们主要参考了以下网络：

• Restormer
• Uformer
• SENet

在以上网络的基础上，我们对网络做出了一定的改进，并提出了一些自己的创新方法：

• Restormer_SE：加入了channel attention和SE module，其中channel attention和MDTA并行形成Hybrid Attention

• Restormer_KA：加入了利用Dynamic Conv的Kernel Attention，有NAN问题

• Restormer_KAv2：加入了利用普通卷积的Kernel Attention，没有nan问题；Kernel Attention分窗大小变化（16->8->4）

• Restormer_KAv3：加入了利用普通卷积的Kernel Attention，分窗大小不变（16）

• Uformer_KAv1：加入了利用普通卷积的Kernel Attention，在feature map大小为64和32时（输入大小为128）与window_attention形成两个分支，分窗大小不变均为16；而且仅在未进行shift window attention时进行kernel attention

• Uformer_KAv2：基于v1，更新了Kernel Attention的结构

• Uformer_KAv3：

• Uformer_KAv4：串行结构代码，窗口卷积+global kernel

• Uformer_KAv5：开始使用并行结构代码

• Uformer_KAv6：BasicUformerLayer层数增加，只有LeWinTransformerBlock去除了shift window，LeWinTransformerBlock没有去除了shift window

• Uformer_KAv7：只使用global kernel

• Uformer_KAv8：BasicUformerLayer层数没有增加，LeWinTransformerBlock_KA和LeWinTransformerBlock都去除了shift window

• Uformer_KAv9：基于KAv7，去除了shift window

• Uformer_KAv10：基于KAv4，去除了SA和SE

• Uformer_KAv11：基于KAv4，去除了global kernel的se

• Uformer_KAv12：基于KAv4，去除了移位

• Uformer_KAv13：基于KAv4，普通卷积改为深度可分离卷积

Note：Restormerv2和Restormerv3的代码应该有问题，在ConvLayer部分权重的维度变换有问题

Results and Analysis

统计学习

我们设置学习率为1e-4，使用AdamW优化器，以L1Loss和SSIMloss作为损失函数。

以下为1000epoch训练下，各个模型的效果对比：

模型	PSNR	SSIM	Params	FLOPs
Restormer	38.74	0.97994
Restormer_SE	38.41	0.97752
Restormer_KAv2	-	-
Uformer(bs=8)	38.49	0.97810	20.628M	10.308G
Uformer	38.76	0.97927	20.628M	10.308G
Uformer_KAv1	38.63	0.97798	21.953M	10.886G
Uformer_KAv2	38.78	0.97920	24.196M	12.704G
Uformer_KAv3	38.78	0.97940	24.196M	12.704G
Uformer_KAv4	38.82	0.97947	24.667M	13.548G

除此之外，我们还测试了一些模型在不同epoch训练下的效果：

模型	epoch	PSNR	SSIM	模型	epoch	PSNR	SSIM
Restormer	500	37.30	0.97369	Restormer_SE	500	37.46	0.97369
	1000	38.76	0.97927		1000	38.41	0.97752
	2000	39.47	0.98284		2000	39.35	0.98114
	3000	39.79	0.98382		3000	39.64	0.98200

模型	epoch	PSNR	SSIM	模型	epoch	PSNR	SSIM
Uformer	1000	37.30	0.97369	Uformer_KAv4	1000	38.82	0.97947
	2000	39.60	0.98264		2000	39.36	0.98219

更多的，我们在Rain100H上取200对图片作为训练集，经过训练，测试结果如下：

模型	epoch	PSNR	SSIM	模型	epoch	PSNR	SSIM
Uformer	1000	35.06	0.95970	Uformer_KAv4	1000	29.62	0.86394

tiny paper

1000epoch情况下：

模型	PSNR	SSIM	Params	FLOPs
Uformer bs=6	38.67999	0.97879	20.628M	10.308G
Uformer bs=12	37.87080	0.97472	20.628M	10.308G
Uformer_tinypaperv1	38.00910	0.97563
Uformer_tinypaperv2	38.29627	0.97658
Uformer_KAv4	38.81872	0.97947	24.667M	13.548G
Uformer_KAv5	38.72450	0.97846	24.854M	17.238G
Uformer_KAv7	38.71809	0.97864
Uformer_KAv8	38.78743	0.97921
Uformer_KAv9	38.81705	0.97939
Uformer_KAv10	38.78690	0.97930
Uformer_KAv11	38.77819	0.97959
Uformer_KAv12	38.67999	0.97879
Uformer_KAv13	38.75458	0.97909	20.416M	9.82G

ssim: , psnr: 2000epoch情况下：

模型	PSNR	SSIM	Params	FLOPs
Uformer	39.60271	0.98264	20.628M	10.308G

3000epoch情况下：

模型	PSNR	SSIM	Params	FLOPs
Uformer	39.92405	0.98404	20.628M	10.308G

Training and Evaluation

Training and Testing for Deraining:

Derain	Dataset	Visual Results
Rain100L	Link	Empty
Rain100H	Link	Empty

Note

• Download datasets and put it with the following format.

• We introduce more methods based on Restormer to the single image deraing task.

• The project is based on MMCV, but you needn't to install it and master MMCV. More importantly, it can be more easy to introduce more methods.

• We have modified part of the code for reading data set, specifally for Rain100L, which can be read as the same way with Rain200L. But you have to operate on the dataset. Target and rainy figure should be operated into one figure. You can finish the operation above through dataset_rename.py, but be careful to use them.

|-$ROOT/datasets
├── Rain100L
│   ├── train_c
│   │   ├── norain-001.png
│   │   ├── ...
│   ├── test_c
│   │   │   ├── norain-001.png
│   │   │   ├── ...