[Pytorch] Super-Resolution using Multi-Step Reinforcement Learning

Implementation of Super-Resolution model using Reinforcement learning based on Multi-Step Reinforcement Learning for Single Image Super-Resolution paper with Pytorch.

Contents

• Introduction
• Train
• Test
• Demo
• Evaluate
• References

Introduction

This implementation uses PixelRL as the core of reinforement learning model, I use my 4 available Super-Resolution models for actions instead of EDSR and ESRGAN as in paper. Because changing the actions are easy so you can use other super-resolution models as your actions. I ignore patch-wise agent and keep t_max=5 as in PixelRL paper.

Index	Action
0	pixel value -= 1
1	do nothing
2	pixel value += 1
3	ESPCN
4	SRCNN
5	VDSR
6	FSRCNN

The actions.

Train

Dataset:

• Train: T91 + General100 + BSD200
• Validation: Set14
• Test: Set5

You run this command to begin the training:

python train.py --scale=2              \
                --steps=2000           \
                --batch-size=64        \
                --save-every=50        \
                --save-log=0           \
                --ckpt-dir="checkpoint/x2"

• --save-log: if it's equal to 1, train loss, train rewards, train metrics, validation rewards, validation metrics will be saved every save-every steps.

NOTE: if you want to re-train a new model, you should delete all files in checkpoint sub-directory. Your checkpoint will be saved when above command finishs and can be used for the next times, so you can train a model on Google Colab without taking care of GPU time limit.

I trained the models on Google Colab in 2000 steps:

You can get the models here:

• PixelRL_SR-x2.pt
• PixelRL_SR-x3.pt
• PixelRL_SR-x4.pt

The log informations of my training process are ploted below, these plots are smoothed by Exponential Moving Average with alpha=0.2:

Loss (Left), reward (center), PSNR (right) of scale x2.


Loss (Left), reward (center), PSNR (right) of scale x3.


Loss (Left), reward (center), PSNR (right) of scale x4.

Test

I use Set5 as the test set. After Training, you can test models with scale factors x2, x3, x4, the result is the average PSNR of all images.

python test.py --scale=2 --ckpt-path="default"

--ckpt-path="default" mean you are using default model path, aka checkpoint/x{scale}/PixelRL_SR-x{scale}.pt. If you want to use your trained model, you can pass yours to --ckpt-path.

Demo

After Training, you can test models with this command, the result is the sr.png.

python demo.py --scale=2             \
               --ckpt-path="default" \
               --draw-action-map=0   \
               --image-path="dataset/test1.png"

--draw-action-map: If it's equal to 1, an action map will be save to action_maps directory every step.

--ckpt-path is the same as in Test

Evaluate

I evaluated models with Set5, Set14, BSD100 and Urban100 dataset by PSNR, the Bold texts are the best results:

Dataset	Scale	Bicubic	ESPCN	SRCNN	FSRCNN	VDSR	PixelRL-SR
Set5	2	32.0500	38.2830	36.7996	38.7593	36.9849	38.7633
	3	28.8415	34.6919	34.2977	34.5420	34.2582	34.6914
	4	26.8905	32.0646	32.1393	31.9589	31.9323	32.0646
Set14	2	28.5027	34.4974	33.4307	34.5086	33.3692	34.4900
	3	25.8909	31.3246	31.4633	31.2409	31.0208	31.3248
	4	24.3709	29.2934	29.6675	29.3272	29.3366	29.3933
BSD100	2	28.3979	34.3377	33.3674	34.4503	33.4341	34.4501
	3	25.9977	31.3524	31.1648	31.2723	32.0901	31.3525
	4	24.7431	29.7331	29.6832	29.6845	29.6939	29.7331
Urban100	2	25.3959	31.6791	30.2185	31.6858	30.5529	31.6963
	3	X	X	X	X	X	X
	4	21.8013	27.0724	26.9614	27.0038	26.9200	27.0724

Butterfly.png in Set5, Bicubic (left), PixelRL-SR x2 (center), High Resolution (right).

Baboon.png in Set14, Bicubic (left), PixelRL-SR x2 (center), High Resolution (right).

img_001.png in Urban100, Bicubic (left), PixelRL-SR x2 (center), High Resolution (right).

References

• Multi-Step Reinforcement Learning for Single Image Super-Resolution: https://ieeexplore.ieee.org/document/9150927
• Fully Convolutional Network with Reinforcement Learning for Image Processing: https://arxiv.org/abs/1912.07190
• Real-Time Single Image and Video Super-Resolution Using an Efficient Sub-Pixel Convolutional Neural Network: https://arxiv.org/abs/1609.05158
• Image Super-Resolution Using Deep Convolutional Networks: https://arxiv.org/abs/1501.00092
• Accurate Image Super-Resolution Using Very Deep Convolutional Networks: https://arxiv.org/abs/1511.04587
• Accelerating the Super-Resolution Convolutional Neural Network: https://arxiv.org/abs/1608.00367
• FightingSrain/Pytorch-pixelRL: https://github.com/FightingSrain/Pytorch-pixelRL
• T91, General100, BSD200: http://vllab.ucmerced.edu/wlai24/LapSRN/results/SR_training_datasets.zip
• Set5: https://filebox.ece.vt.edu/~jbhuang/project/selfexsr/Set5_SR.zip
• Set14: https://filebox.ece.vt.edu/~jbhuang/project/selfexsr/Set14_SR.zip
• BSD100: https://filebox.ece.vt.edu/~jbhuang/project/selfexsr/BSD100_SR.zip
• Urban100: https://filebox.ece.vt.edu/~jbhuang/project/selfexsr/Urban100_SR.zip