WSL2, AMD GPU, DirectML 환경에서 super-resolution 모델 (BSRGAN) 돌려보기

시스템 사양 (환경)

CPU - 5900X

Memory - 96GB

Mainboard - Gigabyte, X570 Aorus Master

GPU - RX 6900XT (reference)

WSL - Ubuntu 20.04

 

WSL2 환경에서 AMD GPU와 (6900XT) DirectML을 활용하여 Super-resoltion 모델을 돌려보았다.

Super-resolution 모델이란 닳고 닳은 이미지를 딥러닝 기술을 활용하여 새 이미지처럼 만들어 주는 것이다.

 

 

많은 모델이 개발되어 공유되고 있지만 AMD GPU와 DirectML이라는 제한된 환경 탓에

패키지 버전 이슈가 없는 모델로 찾고 찾아 BSRGAN 이라는 모델을 선택하였다.

모델에 대한 자세한 설명은 아래 깃헙 링크에 있다.

 

GitHub - cszn/BSRGAN: Designing a Practical Degradation Model for Deep Blind Image Super-Resolution (ICCV, 2021) (PyTorch) - We released the training code!

GitHub - cszn/BSRGAN: Designing a Practical Degradation Model for Deep Blind Image Super-Resolution (ICCV, 2021) (PyTorch) - We

 

또한 참고한 코랩 튜토리얼은 아래와 같다.

 

SwinIR demo on Real-World Image SR.ipynb - Colaboratory (google.com)

Google Colaboratory Notebook

 

먼저 conda 가상환경을 만들자

1. conda create -n swin python=3.8 입력

환경이름이 swin인 이유는 원래 SwinIR이라는 모델을 돌려보고자 했으나 버전 문제로 실패하였고

BSRGAN을 시도 했으나 이름 바꾸는걸 깜빡했다.

 

 

2. git clone https://github.com/cszn/BSRGAN.git 을 입력하여 깃헙 자료를 받자

 

 

3. sudo apt install libblas3 libomp5 liblapack3 을 입력하여 모델 사용에 필요한 패키지들 설치

 

 

4. 다운받은 깃헙 자료 폴더에 위치해서 (BSRGAN 폴더 안에서)
wget https://github.com/cszn/KAIR/releases/download/v1.0/BSRGAN.pth -p model_zoo 을 입력하여 모델 다운로드

 

 

5. conda activate swin 을 입력하여 가상환경을 활성화 한 후,

pip install torchvision==0.9.0 을 입력하여 torchvision을 먼저 설치한다.

공식 가이드에는 pytorch-directml을 설치하는 것으로 안내하나,
가이드대로 설치하면 호환되지 않는 torchvision 버전이 설치된다. 따라서 먼저 torchvision을 설치한다. 

(관련 링크, How to get available devices and set a specific device in Pytorch-DML? · Issue #165 · microsoft/DirectML · GitHub)

 

 

6. pip install pytorch-directml 을 입력하여 pytorch 설치

 

 

7. tqdm, opencv-python, matplotlib 설치

각각 또는 한꺼번에 pip install 패키지 를 해준다.

 

 

8. 깃헙으로 받은 BSRGAN 디렉토리 내부에서 sudo nano main_test_bsrgan.py 를 입력하여 코드를 수정하자.

기존 코드는 cuda를 이용하는 방식으로 작성되어 있으므로 해당 부분들을 바꾸거나 주석처리 해야한다.

Device 정의부분에 'dml' 입력해준다.

 

 

이외에 cuda를 포함하는 코드들이 보일텐데 (torch.cuda.empty_cache()) 모두 다 주석처리 해준다.

 

 

9. 수정이 되었다면 python main_test_bsrgan.py 를 입력하여 모델을 돌려보자.

돌리면 아래와 같이 알아서 돌아갈텐데, BSRGAN 폴더 안의 testsets/RealSRSet 에 있는 사진들을 처리한다.

처리된 사진은 RealSRSet_results_x4 폴더에 저장된다.

 

 

성능을 확인해 보고자 구글링해서 찾은 아래에 보이는 저화질의 고양이 사진을 넣어보았다.

 

 

결과는

 

 

끝.