SR Trainer

• This repository helps to train and test various deep learning based super-resolution models by changing few configurations(.yaml)

Contents

• Updates
• Models
• Features
• Usage
• Citation

Updates

• News (2024-03-08): Fixed minor issues and being preparing to implement MLFlow
• News (2023-12-08): Implemented warm start and finetune strategy for training
• News (2023-08-21): Implemented profiling model to check its FLOPs, inference time, memory footprint, etc
• News (2023-05-12): Implemented conversion for ONNX and TensorRT inference
• News (2023-05-01): Refactored codes
• News (2023-02-13): Implemented profiler to check spec of models
• News (2023-01-10): Implemented Wandb to monitor status of training
• News (2023-01-01): Implemented Knowledge Distillation
• News (2022-12-30): Refactoring codes for tester and metricer
• News (2022-12-30): Refactored codes for train
• News (2022-12-22): Implemented interlace degradation
• News (2022-08-25): Implemented calculator for quantitative score
• News (2022-07-23): Implemented multi-gpus train approach using pytorch distributed data parallel
• News (2022-07-05): Implemented test codes for video and image
• News (2022-07-02): Implemented train codes for super-resolution models such as EDSR, BSRGAN, RealESRGAN, SwinIR and SCUNet

Models

• BSRGAN
• EDSR
• Real-ESRGAN
• SCUNET
• SwinIR

Features

• Training super-resolution using NET, GAN, Knowledge Distillation approaches
• Testing super-resolution models for image and video
• Converting PyTorch based super-resolution models to ONNX

Train Approaches

• NET
• GAN
• WS
• FT

SR Trainer Structures

1. Configs to manage hyperparameters for train

configs/ ├── models │ ├── BSRGAN.yaml │ ├── EDSR.yaml │ ├── RealESRGAN.yaml │ └── SCUNET.yaml │ └── SwinIR.yaml ├── test │ ├── image.yaml │ └── video.yaml ├── train │ ├── GAN.yaml │ └── NET.yaml │ └── WS.yaml │ └── KD.yaml ├── test.yaml └── train.yaml

2. Scripts for train and test

├── archs │ ├── BSRGAN │ │ ├── README.md │ │ └── models.py │ ├── EDSR │ │ ├── README.md │ │ └── models.py │ ├── RealESRGAN │ │ ├── README.md │ │ └── models.py │ ├── SCUNet │ │ ├── README.md │ │ └── models.py │ ├── SwinIR │ │ └── models.py │ ├── Utils │ │ ├── __init__.py │ │ ├── blocks.py │ │ └── utils.py │ ├── __init__.py ├── data │ ├── __init__.py │ ├── dataset.py │ └── degradation.py ├── learning_rate_scheduler │ ├── __init__.py ├── loss │ ├── __init__.py │ ├── gan_loss.py │ ├── l1_charbonnier_loss.py │ └── perceptual_loss.py ├── metric │ └── metrics.py ├── metricer.py ├── models.py ├── optim │ ├── __init__.py ├── tester.py ├── trainer.py └── utils.py

How to use SR Trainer

1. Modifying train.yaml in the configs directory with your own hyperparamters to train

Train

```yaml

train.yaml

A user should change modelname and traintype to train

hydra: run: dir: ./outputs/train/${now:%Y-%m-%d}/${now:%H-%M-%S}

defaults: - self - train: traintype # PSNR, GAN - models: modelname # EDSR ```

```yaml

PSNR or GAN.yaml

A user should specify train directory in datasets

dataset: train: type: ImagePairDegradationDataset hrdir: datasetspath ```

```yaml

model.yaml

A user should specify the path of checkpoint in order to resume your train

generator: path: "" # /model_zoo/edsr.pth

discriminator: path: "" # /model_zoo/discriminator.pth ```

Test

```yaml

test.yaml

A user should change modelname and testtype to test with various models

hydra: run: dir: ./outputs/test/${now:%Y-%m-%d}/${now:%H-%M-%S}

defaults: - self - test: testtype # image, video - models: modelname # EDSR ```

```yaml

model.yaml

A user should specify the path of pre-trained to load weights, in order to inference your model

generator: path: "" # /model_zoo/edsr.pth ```

Valid

```yaml

valid.yaml

A user should change modelname and vaildtype to valid with various models

hydra: run: dir: ./outputs/SCUNET/valid/${now:%Y-%m-%d}/${now:%H-%M-%S}

defaults: - self - valid: validtypr # quantitative - models: modelname # EDSR ```

```yaml

quantitative.yaml

A user should specify the path of hrdir & lrdir to load custom dataset and metrics to valid your models

savepath: "quantitative" dataset: hrdir: "/workspace/SuperResolution/example/Ref" lr_dir: "/workspace/SuperResolution/example/Dis"

metrics: [psnr, ssim, lpips, erqa] ```

Profiling

```yaml hydra: run: dir: ./outputs/EDSR/profile/${now:%Y-%m-%d}/${now:%H-%M-%S}

defaults: - self - models: EDSR

data: batch: 1 channel: 3 width: 1920 height: 1088 ```

ONNX Converting

```yaml hydra: run: dir: ./outputs/EDSR/onnx/${now:%Y-%m-%d}/${now:%H-%M-%S}

defaults: - self - models: EDSR

data: dynamic_shape: True batch: 1 channel: 3 width: 1920 height: 1080 ```

3. Run

```python

trainer.py for train a model using single GPU

CUDAVISIBLEDEVICES=0 python trainer.py

trainer.py for train a model using multiple GPUs

CUDAVISIBLEDEVICES=0,1,2,3 python trainer.py ```

```python

tester.py for test a model

CUDAVISIBLEDEVICES=0 python tester.py ```

```python

pytorch2onnx.py for converting a model to ONNX

CUDAVISIBLEDEVICES=0 python pytorch2onnx.py ```

```python

profiler.py for profiler a model

CUDAVISIBLEDEVICES=0 python profiler.py ```

Citation

If the repository helps your research or work, please consider citing AutomaticSR.
The following is a BibTeX reference. The BibTeX entry requires the url LaTeX package.

latex @misc{chung2022AutomaticSR, author = {Heaseo Chung}, title = {{AutomaticSR}: Open source for various super-resolution trainer and tester}, howpublished = {\url{https://github.com/HeaseoChung/Super-resolution}}, year = {2022} }