421zuoduan / WIM-code

Geek Repo:Geek Repo

Github PK Tool:Github PK Tool

A Novel Window-Interaction Module Based on W-MSA

Abstract: W-MSA proposed by Swin Transformer has limitations in facilitating information interaction between windows. To address this, we introduce a module that utilizes convolution to achieve inter-window information interaction across different regions. Experiments demonstrate that our proposed module, when combined with W-MSA in a dual-branch structure, outperforms the simple W-MSA. In the deraining task conducted on the Uformer, we observe a 0.14dB improvement in performance.


This repository contains the code used to conduct the experiments for our ICLR 2024 Tiny Paper submission for the paper titled "A Novel Window-Interaction Module Based on W-MSA".

Our Work

The final version can be found in models/compared_trans/Uformer_KAv4/model.py, and model with fewer params can be found in models/compared_trans/Uformer_KAv13/model.py. Following codes are for ablation study.

  • Uformer_KAv7: without window convolutions, add global convolution only
  • Uformer_KAv9: based on KAv7, remove the shift window in W-MSA
  • Uformer_KAv10: based on KAv4,remove the operation of Self Attention and SE module
  • Uformer_KAv11: based on KAv4, remove the SE module for global convolution kernel
  • Uformer_KAv12: based on KAv4, remove the shift window in W-MSA

Datasets

The dataset employed in this study is Rain100L(comprising 200 images for training and 100 images pairs for testing) and Rain200L(comprising 1800 images for training and 200 images pairs for testing). The dataset originally provided at the disclosed URL has been replaced with Rain200L. To ensure the dataset is in the correct structure, dataset_process.py can be utilized for preprocessing.

Derain Dataset
Rain100L Link
Rain200L Link

Training Setting

The learning rate is typically set to 1e-4, and the AdamW optimizer is employed with a random seed set to 1. The input image size for the model is 128*128, and both L1 loss and SSIM loss are utilized as loss functions. We have trained models for 1000 epochs. Details can be found in configs/option_model.py

Training and Evaluation

  • Check configs/option_model.py for configuration.
  • Training. Revise and run python run_derain.py for training
  • Evaluation. Revise and run python run_derain_test.py for testing.

Results

Evaluation of Models

Model PSNR SSIM Params FLOPs
Uformer 38.680 0.9788 20.628M 10.308G
Ours-L 38.820 0.9795 24.667M 13.548G
Ours-T 38.755 0.9791 20.416M 9.82G

Here Ours-T denotes model with less params, Ours-L denotes model with more params. The parameters of the model will be released later.

Note

  • Download datasets and put it with the following format.

  • 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 dataset, 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_process.py, but be careful to use them.

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

About


Languages

Language:Python 74.3%Language:C++ 16.8%Language:Cuda 8.7%Language:Shell 0.1%Language:C 0.1%Language:Batchfile 0.0%Language:Makefile 0.0%Language:CSS 0.0%Language:Dockerfile 0.0%