저번 포스트에 이어 SR(Super Resolution) task를 위해 real-ESRGAN을 사용하는 방법을 소개하려고 한다. 이 모델은 basicsr 과 realesrgan 라이브러리를 사용하여 학습하기 때문에 기존 pytorch기반의 베이스라인 주피터노트북 파일로 진행하지 않고, 프로그램 작동하듯이 파이썬 파일로 진행하게 된다.
데이콘 대회 : AI양재허브 인공지능 오픈소스 경진대회
1등팀 코드 : [Private 1위 25.00327] RRDBNet + Geometric Self-Ensemble
--실험 환경--
Main
OS : Ubuntu 20.04.5 LTS
CPU : Intel(R) Xeon(R) Gold 6230R CPU @ 2.10GHz
GPU : A5000 * 4
Sub
Colab Pro+
(아마도 학습을 main에서 돌리고, 결과확인을 Colab으로 한 듯하다.)
우선 학습 환경은 최소 V100그래픽카드는 필요하다. 1등팀 환경은 A5000 x 4였으니 24GB x 4 = 96GB GPU 메모리 환경에서 905000iter 이상을 돌렸다. V100(32GB)에서 20만 iter를 돌리는데 3일이 걸렸으니 학습을 마치려면 A5000 x 4에서도 최소 4일이상 걸린다는 계산이 나온다.
git clone https://github.com/Jinwoo1126/Dacon_SR.git
cd Dacon_SR
터미널 환경에서 git을 clone해온다. 더 하드코어하게 real-ESRGAN을 밑바닥에서 부터 컨트롤해보고 싶다면 여기 real-ESRGAN github를 참조바란다. 1등팀도 여기를 기준으로 실험을 설계해 나갔다.
# Install basicsr - https://github.com/xinntao/BasicSR
# We use BasicSR for both training and inference
!pip install basicsr
# facexlib and gfpgan are for face enhancement
!pip install facexlib
!pip install gdown
!pip install gfpgan
!pip install -r requirements.txt
!python setup.py develop
필요한 라이브러리와 실험환경을 설정한다.
## can download .pth file manually at https://drive.google.com/file/d/1piw_MOIE5bTH3-o9rmWqp3uIZoYcc5Wl/view
gdown 1piw_MOIE5bTH3-o9rmWqp3uIZoYcc5Wl #get pth file
mv net_g_905000.pth weights #move pth file to weights folder
pretrained model을 다운 받는 방법이다. 여기서는 이 파일을 학습으로 만들어 내야한다. 905000은 돌린 iteration 숫자이다.
## can download compressed img file manually at https://drive.google.com/file/d/17Ui8Pc6NiPTd6dBsZJr8XS4l9S9F-wyL/view
gdown 17Ui8Pc6NiPTd6dBsZJr8XS4l9S9F-wyL #get compressed img file(.zip)
mv open.zip inputs #move zip file to inpus folder
unzip inputs/open.zip -d inputs/ #unzip zip file
학습시킬 이미지 데이터를 다운 받는 방법인데, clone한 레포에 train이미지들만 수동으로 넣어줘도 된다.
폴더 구조는 다음과 같다.
./Dacon_SR
├── ...
├── inputs
│ ├── train # 이미지 데이터 넣어줘야 함
| | ├── hr
│ │ | ├── 0000.png
│ │ | ├── 0001.png
│ │ | ├── 0002.png
| | | ...
│ │ ├── lr
| | ├── 0000.png
| | ├── 0001.png
| | ├── 0002.png
| | ...
│ ├── test # 이미 있음
│ ├── 20000.png
│ ├── 20001.png
│ ├── ...
├── options
│ ├── ...
│ ├── finetune_realesrnet_x4plus_pairdata.yml # argparser 고칠 yaml 파일
├── realesrgan
│ ├── ...
│ ├── train.py
├── ...
├── results # inference한 파일이 저장됨, inference 시 생성
├── scripts
├── submission # results의 파일들 zip파일로 묶음
├── weights
│ ├── net_g_905000.pth # 다운받아서 넣어줘야 함
├── ...
├── inference_rrdbnetrot.py # inference 실행 파일
추가적으로 real-ESRGAN 제작자가 pretrain시켜놓은 모델이 있는데 다운받아서 적용하는걸 추천했다.
wget https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth -P weights
Download pre-trained models: RealESRGAN_x4plus.pth
weights
폴더나 experiments/pretrained_models/
를 만들어서 넣어놓으면 된다.
라이브러리를 불러와서 학습하는 방법이기 때문에 yaml파일을 수정해서 학습시키게된다. ./options/finetune_realesrnet_x4plus_pairdata.yml
파일을 수정한다.
# general settings
name: finetune_RealESRNetx4plus_400k_pairdata
model_type: RealESRNetModel
scale: 4 # 4배로 화질 상승!
num_gpu: auto
manual_seed: 0
# USM the ground-truth
# 원본 사진 사용 여부이다.
l1_gt_usm: True
percep_gt_usm: True
gan_gt_usm: False
high_order_degradation: False # do not use the high-order degradation generation process
# dataset and data loader settings
datasets:
train: # the 1st test dataset
name: Dacon_train
type: PairedImageDataset
# 만약 학습 이미지 데이터 경로를 바꾸고 싶으면 여기서 수정한다.
dataroot_gt: ./inputs/train/hr
dataroot_lq: ./inputs/train/lr
io_backend:
type: disk
# 파일 이름을 메모리에 올리는 방식을 선택하게 된다.
# 자세한건 ./realesrgan/data/realesrgan_paired_dataset.py파일 참조
gt_size: 512 # lr의 이미지 사이즈
use_hflip: true # Augmentation : hflip 수행
use_rot: true # Augmentation : rotation 수행
# data loader
use_shuffle: true
num_worker_per_gpu: 4
# 가끔 num_worker때문에 오류나는 가상환경들이 있는데 0으로 설정하면 오류 해결
batch_size_per_gpu: 4
dataset_enlarge_ratio: 1
prefetch_mode: ~
# network structures
network_g:
type: RRDBNet
num_in_ch: 3
num_out_ch: 3
num_feat: 64 # feature 갯수 높일수록 성능 향상 기대
num_block: 23
num_grow_ch: 32
network_d:
type: UNetDiscriminatorSN
num_in_ch: 3
num_feat: 64
skip_connection: True
# path
path:
# use the pre-trained Real-ESRNet model
# 제작자의 pre-trained 모델을 쓰려면 첫번째꺼를 쓰고
# 아무것도 없이 학습한다면 pretrain_network_g: 뒤를 비워주자.
#pretrain_network_g: experiments/pretrained_models/RealESRGAN_x4plus.pth
pretrain_network_g: ./weights/net_g_905000.pth
param_key_g: params_ema
strict_load_g: false
resume_state: ~
# training settings
train:
ema_decay: 0.999
# Optimizer (Adam Optimizer 사용, Learning rate = 0.0004, beta값은 [0.9, 0.99])
optim_g:
type: Adam
lr: !!float 1e-4
weight_decay: 0
betas: [0.9, 0.99]
# Scheduler (200,000iter마다 gamma 값이 0.5씩 줄어들도록 decay 설정)
scheduler:
type: MultiStepLR
milestones: [200000]
gamma: 0.5
total_iter: 999999999
warmup_iter: -1 # no warm up
# losses
pixel_opt:
type: MSELoss
loss_weight: 1.0
reduction: mean
# validation settings
val:
val_freq: !!float 5e3 # 5000iter에 한번씩 validation한다.
save_img: false
pbar: False
metrics:
psnr:
type: calculate_psnr # 평가방식 : PSNR (Peak Signal-to-Noise Ratio)
crop_border: 4
test_y_channel: true
better: higher # the higher, the better. Default: higher
ssim:
type: calculate_ssim
crop_border: 4
test_y_channel: true
better: higher # the higher, the better. Default: higher
# logging settings
logger:
print_freq: 100
# 100iter에 한번 진행상황이 프린트 된다.
save_checkpoint_freq: !!float 5e3
# 모델 저장 주기를 수정할 수 있다.
# 모델은 ./experiments/[your experiments name(default='finetune_RealESRNetx4plus_400k_pairdata')]/models에
# 저장된다.
use_tb_logger: true
wandb:
project: ~
resume_id: ~
# dist training settings
dist_params:
backend: nccl
port: 29500
이 블록 안에도 작성하였지만 실험 상황을 수정할 방법은 다음과 같다.
python realesrgan/train.py -opt options/finetune_realesrnet_x4plus_pairdata.yml
파일 경로만 다 맞고, 라이브러리가 잘 설치되었다면 잘 돌아갈 것이다. wandb를 설정하면 실시간으로 성능도 확인 가능하다.
모델은
./experiments/[your experiments name(default='finetune_RealESRNetx4plus_400k_pairdata')]/models
에 저장되므로 colab 이나 다른 컴퓨터로 계속 상황을 보면서 잘된 모델을weights
폴더로 옮겨주면 된다.이후 inference 하는 과정은 이전 포스트에 적어놨으니 참고바라고 좋은 화질의 사진을 얻길 바란다!!