A PyTorch implementation of Restormer based on CVPR 2022 paper Restormer: Efficient Transformer for High-Resolution Image Restoration.

• Anaconda

• PyTorch

conda install pytorch=1.10.2 torchvision cudatoolkit -c pytorch

Rain100L and Rain100H are used, download these datasets and make sure the directory like this:

|-- data     
    |-- rain100L
        |-- train
            |-- rain
                norain-1.png
                ...
            `-- norain
                norain-1.png
                ...
        `-- test                                                        
    |-- rain100H
        same as rain100L

You can easily train and test the model by running the script below. If you want to try other options, please refer to utils.py.

python main.py --data_name rain100L --seed 0

python main.py --data_name rain100H --model_file result/rain100H.pth

The models are trained on one NVIDIA RTX A6000 GPU (48G). num_iter is 30,000, seed is 1 and milestone is [9200, 15600, 20400, 24000, 27600], the other hyper-parameters are the default values.

Due to the huge demand for GPU memory, we have to reduce the batch_size and patch_size:

Ours: batch_size is [64, 40, 32, 16, 8, 8] and patch_size is [32, 40, 48, 64, 80, 96];

Ours*: batch_size is [32, 20, 16, 8, 4, 4] and patch_size is [64, 80, 96, 128, 160, 192].

More results could be downloaded from MEGA. Here we give some examples for Ours*.