Skip to content

Latest commit

 

History

History
45 lines (27 loc) · 1.3 KB

README.md

File metadata and controls

45 lines (27 loc) · 1.3 KB

Pusnet

This repo is the official code for

Dependencies and Installation

  • Python 3.8.13, PyTorch = 1.11.0

  • Run the following commands in your terminal:

    conda env create -f env.yml

    conda activate pyt_env

Get Started

Training

  1. Change the code in config.py

    line4: mode = 'train'

    `line14: train_data_dir=''

    `line15: test_data_dir=''

  2. Run python pusnet.py

  • Trained models will be saved in 'model_zoo' folder.

Testing

  1. Change the code in config.py

    line4: mode = 'test'

    line36-41: test_pusnet_path = ''

  2. Run python pusnet.py

  • Testing results will be saved in 'results' folder.
  • Here, we provide trained models.

Dataset

  • We train the PUSNet on the DIV2K training dataset, and test it on three testing datasets, including the DIV2K test dataset, 1000 images randomly selected from the ImageNet test dataset

Others

  • The batch_size in config.py should be at least 2*number of gpus and it should be divisible by number of gpus.