/Uniwin

PyTorch implementation of Uniwin("Image Super-resolution with Unified Window Attention".

Primary LanguagePython

Image Super-Resolution with Unified Window Attention

Gunhee Cho and YongSuk Choi

Artificial Intelligence Lab, Hanyang University, Seoul, Korea


This repository is the official Pytorch implementation of Image Super-resolution with Unified Window Attention.

Architecture

arch

Results

Quantitative Results

Qualitative Results

Ablations

Environment

  • Ubuntu 20.04 LTS
  • 4 NVIDIA RTX A5000

Install

pip3 install -r requirements.txt

Preparation

  • Download train dataset (DF2K/ImageNet)
  • Download test dataset (Set5/Set14/BSD100/Urban100/Manga109)

Training

To pretrain with ImageNet data (x2/x3/x4)

torchrun --standalone --nproc_per_node=4 train.py --opt options/train_Uniwin_SRx2_ImageNet_from_scratch.json

torchrun --standalone --nproc_per_node=4 train.py --opt options/train_Uniwin_SRx3_ImageNet_from_scratch.json

torchrun --standalone --nproc_per_node=4 train.py --opt options/train_Uniwin_SRx4_ImageNet_from_scratch.json

To finetune with DF2K data (x2/x3/x4)

torchrun --standalone --nproc_per_node=4 train.py --opt options/train_Uniwin_SRx2_finetune_from_ImageNet_pretrain.json

torchrun --standalone --nproc_per_node=4 train.py --opt options/train_Uniwin_SRx3_finetune_from_ImageNet_pretrain.json

torchrun --standalone --nproc_per_node=4 train.py --opt options/train_Uniwin_SRx4_finetune_from_ImageNet_pretrain.json

To finetune from SRx2 (x3/x4)

torchrun --standalone --nproc_per_node=4 train.py --opt options/train_Uniwin_SRx3_finetune_from_SRx2.json

torchrun --standalone --nproc_per_node=4 train.py --opt options/train_Uniwin_SRx4_finetune_from_SRx2.json

To train from scratch with DF2K (x2/x3/x4)

torchrun --standalone --nproc_per_node=4 train.py --opt options/train_Uniwin_SRx2_DF2K_from_scratch.json

torchrun --standalone --nproc_per_node=4 train.py --opt options/train_Uniwin_SRx2_DF2K_from_scratch.json

torchrun --standalone --nproc_per_node=4 train.py --opt options/train_Uniwin_SRx2_DF2K_from_scratch.json