chaofwang / awsrn

PyTorch code for our paper "Lightweight Image Super-Resolution with Adaptive Weighted Learning Network"

Lightweight Image Super-Resolution with Adaptive Weighted Learning Network

Chaofeng Wang, Zhen Li and Jun Shi, "Lightweight Image Super-Resolution with Adaptive Weighted Learning Network", [arXiv]

The code is based on EDSR (PyTorch)

Dependencies

• Python 3.5
• PyTorch >= 0.4.0
• numpy
• skimage
• imageio
• matplotlib
• tqdm

Code

	git clone [email protected]:ChaofWang/AWSRN.git
	cd AWSRN

1. Abstract
2. Architecture
3. Data
4. Train
5. Test
6. Results
7. Q&A
8. TODO
9. Citation
10. Acknowledgements

Abstract

Deep learning has been successfully applied to the single-image super-resolution (SISR) task with great performance in recent years. However, most convolutional neural network based SR models require heavy computation, which limit their real-world applications. In this work, a lightweight SR network, named Adaptive Weighted Super-Resolution Network (AWSRN), is proposed for SISR to address this issue. A novel local fusion block (LFB) is designed in AWSRN for efficient residual learning, which consists of stacked adaptive weighted residual units (AWRU) and a local residual fusion unit (LRFU). Moreover, an adaptive weighted multi-scale (AWMS) module is proposed to make full use of features in reconstruction layer. AWMS consists of several different scale convolutions, and the redundancy scale branch can be removed according to the contribution of adaptive weights in AWMS for lightweight network. The experimental results on the commonly used datasets show that the proposed lightweight AWSRN achieves superior performance on ×2, ×3, ×4, and ×8 scale factors to state-of-the-art methods with similar parameters and computational overhead.

Architecture

​ Network architecture of the proposed AWSRN

​ (a) Basic RU, (b) wRU, (c) AWRU (Ours)

Data

all scale factor(x2,x3,x4,x8) data:

​ training data DIV2K(800 training + 100 validtion images)

​ benchmark data (Set5, Set14, B100, Urban100, Manga109)

can download from Google Drive. and extract the data to ./DATA path

Train

1. Cd to './code/src', run the following commands to train models.

   #=================== AWSRN =================
   #AWSRN x2
   python main.py --model AWSRN --n_resblocks 4 --scale 2  --save AWSRNx2  --epochs 1000 --reset --patch_size 96
   #AWSRN x3
   python main.py --model AWSRN --n_resblocks 4 --scale 3 --save AWSRNx3  --epochs 1000 --reset --patch_size 144
   #AWSRN x4
   python main.py --model AWSRN --n_resblocks 4 --scale 4  --save AWSRNx4  --epochs 1000 --reset --patch_size 192
   #AWSRN x8
   python main.py --model AWSRN --n_resblocks 4 --scale 8  --save AWSRNx8  --epochs 1000 --reset --patch_size 384
   
   #=================== AWSRN-M =================
   #AWSRN-M x2
   python main.py --model AWSRN --n_resblocks 3 --scale 2  --save AWSRN_Mx2  --epochs 1000 --reset --patch_size 96
   #AWSRN-M x3
   python main.py --model AWSRN --n_resblocks 3 --scale 3  --save AWSRN_Mx3  --epochs 1000 --reset --patch_size 144
   #AWSRN-M x4
   python main.py --model AWSRN --n_resblocks 3 --scale 4  --save AWSRN_Mx4  --epochs 1000 --reset --patch_size 192
   
   #=================== AWSRN-S =================
   #AWSRN-S x2
   python main.py --model AWSRN --n_resblocks 1 --scale 2  --save AWSRN_Sx2  --epochs 1000 --reset --patch_size 96
   #AWSRN-S x3
   python main.py --model AWSRN --n_resblocks 1 --scale 3 --save AWSRN_Sx3_retrain  --epochs 1000  --reset --patch_size 144
   #AWSRN-S x4
   python main.py --model AWSRN --n_resblocks 1 --scale 4 --save AWSRN_Sx4  --epochs 1000  --reset --patch_size 192
   
   #=================== AWSRN-SD =================
   #AWSRN-SD x2
   python main.py --model AWSRND --n_resblocks 1 --n_feats 16 --block_feats 128 --scale 2  --save AWSRN_SDx2  --epochs 1000 --reset --patch_size 96 
   #AWSRN-SD x3
   python main.py --model AWSRND --n_resblocks 1 --n_feats 16 --block_feats 128 --scale 3  --save AWSRN_SDx3  --epochs 1000 --reset --patch_size 144 
   #AWSRN-SD x4
   python main.py --model AWSRND --n_resblocks 1 --n_feats 16 --block_feats 128 --scale 4  --save AWSRN_SDx4  --epochs 1000 --reset --patch_size 192 

Test

Quick start

1. Download models and place them in './code/experiment'.

   All the models can be downloaded from Google Drive

2. Cd to './code/src', run the following commands to get result with paper reported.

   #for examples, test on Set5
   #=================== AWSRN =================
   #AWSRN x2
   python main.py --model AWSRN  --n_resblocks 4 --scale 2  --pre_train ../experiment/AWSRNx2/model/model_latest.pt --save AWSRNx2 --test_only  --dir_data ../../DATA/ --data_test Set5
   #AWSRN x3
   python main.py --model AWSRN  --n_resblocks 4 --scale 3  --pre_train ../experiment/AWSRNx3/model/model_latest.pt --save AWSRNx3 --test_only  --dir_data ../../DATA/ --data_test Set5
   #AWSRN x4
   python main.py --model AWSRN  --n_resblocks 4 --scale 4  --pre_train ../experiment/AWSRNx4/model/model_latest.pt --save AWSRNx4 --test_only  --dir_data ../../DATA/ --data_test Set5
   #AWSRN x8
   python main.py --model AWSRN  --n_resblocks 4 --scale 8  --pre_train ../experiment/AWSRNx8/model/model_latest.pt --save AWSRNx8 --test_only  --dir_data ../../DATA/ --data_test Set5
   
   #=================== AWSRN-M =================
   #AWSRN-M x2
   python main.py --model AWSRN  --n_resblocks 3 --scale 2  --pre_train ../experiment/AWSRN_Mx2/model/model_latest.pt --save AWSRN_Mx2 --test_only  --dir_data ../../DATA/ --data_test Set5
   #AWSRN-M x3
   python main.py --model AWSRN  --n_resblocks 3 --scale 3  --pre_train ../experiment/AWSRN_Mx3/model/model_latest.pt --save AWSRN_Mx3 --test_only  --dir_data ../../DATA/ --data_test Set5
   #AWSRN-M x4
   python main.py --model AWSRN  --n_resblocks 3 --scale 4  --pre_train ../experiment/AWSRN_Mx4/model/model_latest.pt --save AWSRN_Mx4 --test_only  --dir_data ../../DATA/ --data_test Set5
   
   #=================== AWSRN-S =================
   #AWSRN-S x2
   python main.py --model AWSRN  --n_resblocks 1 --scale 2  --pre_train ../experiment/AWSRN_Sx2/model/model_latest.pt --save AWSRN_Sx2 --test_only  --dir_data ../../DATA/ --data_test Set5
   #AWSRN-S x3
   python main.py --model AWSRN --n_resblocks 1 --scale 3  --pre_train ../experiment/AWSRN_Sx3/model/model_latest.pt --save AWSRN_Sx3_retrain --test_only  --dir_data ../../DATA/ --data_test Set5
   #AWSRN-S x4
   python main.py --model AWSRN --n_resblocks 1 --scale 4  --pre_train ../experiment/AWSRN_Sx4/model/model_latest.pt --save AWSRN_Sx4 --test_only  --dir_data ../../DATA/ --data_test Set5
   
   #=================== AWSRN-SD =================
   #AWSRN-SD x2
   python main.py --model AWSRND  --n_resblocks 1 --n_feats 16 --block_feats 128 --scale 2  --pre_train ../experiment/AWSRN_SDx2/model/model_latest.pt --save AWSRN_SDx2 --test_only  --dir_data ../../DATA/ --data_test Set5 
   #AWSRN-SD x3
   python main.py --model AWSRND  --n_resblocks 1 --n_feats 16 --block_feats 128 --scale 3  --pre_train ../experiment/AWSRN_SDx3/model/model_latest.pt --save AWSRN_SDx3 --test_only  --dir_data ../../DATA/ --data_test Set5
   #AWSRN-SD x4
   python main.py --model AWSRND  --n_resblocks 1 --n_feats 16 --block_feats 128 --scale 4  --pre_train ../experiment/AWSRN_SDx4/model/model_latest.pt --save AWSRN_SDx4 --test_only  --dir_data ../../DATA/ --data_test Set5
   

NOTE

1. About AWSRN-S x3 pretrainmodel

   Unfortunately, the pre-training model of AWSRN-S x3 was accidentally deleted by me during the process, so I retrained the model, which may be slightly different from the results reported in the paper.

2. Evaluate PSNR and SSIM.

   if you want to evaluate SSIM for each dataset, you need to save images by appending '--save_result' to test commands.

   Then specify the scale and file name in 'Evaluate_PSNR_SSIM.m' . Run 'Evaluate_PSNR_SSIM.m' to obtain PSNR/SSIM values for paper.

   In addition, the PSNR results obtained in Python are the same as in "Evaluate_PSNR_SSIM.m". So you don't need to run this matlab script in order to verify the value of PSNR.Results

Results

Quantitative Results

​ Performance comparison x2 scale on Set5

For more detail, please refer to our paper .

Visual Results

Visual comparisons over four datasets (Set5, Set14, B100 and Urban100) for ×4 scale

TODO

1. Provide our results on benchmark dataset

   Our all models results can download Baidu Cloud code: h39w or Google Drive

2. Train deeper AWRNs with more parameters

3. TBA...

Q&A

1. Q: why have three AWSRN model files in “./src/model” ?

   A: The two files "awsrn.py" and "awsrnd.py" are due to the unclean code styles in my exploration research process. Since the results in the paper are the two model files, they are published. And "awsrnv2.py" is a unified version, you can use this version if you want to retrain all AWSRN models.

Citation

If you find the code helpful in your resarch or work, please cite the following papers.

@article{wang2019lightweight,
  title={Lightweight Image Super-Resolution with Adaptive Weighted Learning Network},
  author={Wang, Chaofeng and Li, Zhen and Shi, Jun},
  journal={arXiv preprint arXiv:1904.02358},
  year={2019
}

Acknowledgements

This code is built on EDSR (PyTorch). We thank the authors for sharing their codes of EDSR PyTorch version.

We also refer to some other work:

wdsr_ntire2018(PyTorch)

RCAN (PyTorch)

CARN (PyTorch)

awsrn's Issues

how do you plot visualization results on different scale branches

how do you plot visualization results on different scale branches

ImportError

I am getting following error while running the training script:
"ImportError: cannot import name '_worker_manager_loop'"
My pytorch version is 1.0.1.

MutiAdds

Hey @ChaofWang,

Thanks for sharing the code. How did you count Mult-Adds of your models?

Thanks,

can't download benchmark dataset results

Thank you for publishing your paper and code!

I'm trying to download your Set5 etc results. But I can't seem to be able to. Chrome/Google translate cannot handle the Baidu website. So I'm trying to make the download work "blind" because I don't understand the language.

Whenever I try to download your results, I get a BaiduNetdisk_6.7.3.exe download instead. I definitely don't want to run an EXE on my PC, for obvious reasons. I just want to download the images. Is that possible somehow?

Thanks!

something wrong with the result in the paper?

Hi, is this value right in your paper?

Some questions about training

Hello, I find that you use 'restart' to train the model, which means 1000 epochs plus 1000 epochs with new learning rate.
Can you tell me why you use 'restart'?
Does 'restart' has some benefits? @ChaofWang

ERROR when testing with option --cpu

hi,

When I run a test and add the option '--cpu' I get an error:

  File "d:\0_PyAPPS\PyTEST\ChaofWang_AWSRN\code\model\awsrn.py", line 134, in forward
    x = (x - self.rgb_mean.cuda()*255)/127.5
RuntimeError: Expected object of type torch.FloatTensor but found type torch.cuda.FloatTensor for argument #3 'other'

MultiAdds

@ChaofWang Hello, could you please tell me your test code of MultAdds, I used the thop library, and the measured data are quite different from those in your paper.Thanks!

DIV2K scale x8数据集

你好~
我在官网上没有看到x8的数据集，想问您一下您做x8的时候用的数据集是自己用matlab处理得到的吗？
在训练scale x8的模型的时候需要对代码中的data部分做一些特殊的修改吗？还是直接将数据集放进去就可以用？
打扰啦

How to compute the flops and params ?

sorry to disturb you , I am interested in the lightweight network too. But I hardly find good codes to compute the flops , params and depth for my network. So can you provide you code about that for references. Thks a lot

About Multi-Adds

Thank you for your work, I want to know that how to calculate the Mult-Adds in your work, and could you provide the code for calculating Mult-Adds.

Mult-Adds

Hi, thank you for your work, I want to know that how to calculate the Mult-Adds in your work.

About Mult-Adds

Hi, thank you for your work, I want to know that how to calculate Mult-Adds in your paper.

Some questions about Data specifications in option.py

I want to train my own data,but i don't know how to place my data in a correct dataset. Do you mean i should place training data and test data in one folder,but the number is different?must I have lr dataset directory?

询问论文细节

1、P7表格中的ours指标是否使用self-ensemble?
2、我使用https://github.com/ceykmc/pytorch_model_summary
脚本，3x16x16作为输入，得到2x的AWSRN模型，理论Madd 712,361,984，换算到论文中的1280x720（*3600）， 是2565G左右，但是论文中是320.5G，相差很远，我们的统计方式是哪里有不同？
3、Table4中，没看懂3579后的那一列数字是指什么，是训练完后的scale值吗？WDSR里面各个分支都有足够多权重，但是AWSRN-B中7x7却几乎零权重，这个是理论上是零，还是巧合？后面表格中MAdd和参数量有没有统计这条支路？可以不统计这条支路，指标可以更好看。

AWMS can be re-paramed

inp=torch.randn(1,16,8,8)
awms=AWMS()#去掉wn的版本
opt1=awms(inp)
k3=awms.tail_k3.weight
b3=awms.tail_k3.bias
k5=awms.tail_k5.weight
b5=awms.tail_k5.bias
k7=awms.tail_k7.weight
b7=awms.tail_k7.bias
k9=awms.tail_k9.weight
b9=awms.tail_k9.bias
scale_k3=awms.scale_k3.scale
scale_k5=awms.scale_k5.scale
scale_k7=awms.scale_k7.scale
scale_k9=awms.scale_k9.scale
conv9r=nn.Conv2d(in_channels=16, out_channels=3, kernel_size=9, stride=1,padding=4, bias=True)
conv9r.weight.data=scale_k3 * torch.nn.functional.pad(k3, [3, 3, 3, 3])+scale_k5 * torch.nn.functional.pad(k5, [2, 2, 2, 2])+scale_k7 * torch.nn.functional.pad(k7, [1, 1, 1, 1])+scale_k9 * k9
conv9r.bias.data=scale_k3 * b3+scale_k5 * b5+scale_k7 * b7+scale_k9 * b9
opt2=awms.pixelshuffle(conv9r(inp))
delta = opt2 - opt1
print(np.abs(delta.detach().cpu().numpy()).mean())

输出1.13682596e-07

4个卷积4次pixelshuffle变成1个9x9卷积1次pixelshuffle
计算量9 * 9/(3 * 3+5 * 5+7 * 7+9 * 9)约为一半
至于是否wn的转化可参考weight_norm.py的代码进行卷积参数的互相转化

compare with WDSR and EDSR

Hi, thanks for the code. I notice that your model is developed by WDSR, would you please provide the comparision results with WDSR and EDSR? Thanks.

https://github.com/JiahuiYu/wdsr_ntire2018/blob/master/wdsr_b.py