[Computer Vision] Super Resolution 구현 - SRGAN

Super Resolution 구현 - SRGAN

안녕하세요, 데이터과학을 전공하고 있는 황경하입니다. 오늘은 저번에 이어 프로젝트로 진행하고 있는 Super Resolution에 대해 이야기해 보겠습니다. 저번에는 SRCNN의 논문을 읽고 직접 구현해 보며 결과를 살펴봤습니다. 그리고, Fine-Tuning을 통해 성능을 개선시킨다는 계획도 밝혔었죠. 하지만, 계획이 바뀌었습니다. 자료를 찾아보니, GAN 모델을 활용한 Super Resolution의 결과가 일반적으로 CNN계열 모델(SRCNN, FSRCNN 등)의 결과보다 정성적인 평가에서 더 좋다고 알려져 있더라고요. 저희는 정해진 기간 동안 빠르게 결과를 내야 하는 프로젝트이기 때문에 CNN 모델들을 과감히 포기하고, GAN 모델 구현을 시작했습니다.

 

따라서, 오늘은 GAN 모델을 처음으로 Super Resolution에 적용한 SRGAN 모델을 구현해 보고, 성능을 비교해 보겠습니다.

 

서론

모든 연구가 그렇듯 말로만 좋다고 하는 것은 의미가 없죠. 그래서 우리는 Accuracy, MSE 등의 Metrics를 이용한 결과를 제시합니다. 이는 숫자로 표현된 결과로 "정량적 평가"라고 합니다. 이와 반대로, "정성적 평가"라는 것은 숫자로 표현되는 것이 아닌 인간이 눈으로 결과를 보고 난 후 내린 결과입니다. 예를 들어, Super Resolution 작업을 거쳐서 고해상도로 복원된 이미지가 원본 이미지와 MSE가 매우 낮다고 합시다. 그러면, 정량적 평가는 매우 좋은 결과를 보이는 것이죠. 하지만, 눈으로 보기에는 결과가 그리 좋지 않을 수 있습니다. 즉, 항상 정량적 평가 = 정성적 평가가 아니라는 것이죠.

Why?

그럼 CNN 계열 모델이 왜 GAN 모델보다 정성적 평가가 좋지 않을까요? 정확한 이유는 이 블로그에 정리해 두었으니, 자세한 내용이 궁금하신 분은 블로그를 참고해 주시면 감사하겠습니다. 여기서는 짧게만 이야기해 볼게요. 그 이유는, Object Function에 있습니다. CNN계열 모델의 Object Function은 MSE였습니다. 이 경우  작동방식은 pixel-wise 방식으로 원본 이미지와 복원 이미지 사이에 픽셀값이 유사해지도록 작동합니다. 그러면, 이런 경우를 살펴보죠. 처음에 복원 이미지는 원본 이미지와 많이 다를 것입니다. 그러면, 모델은 이 두 값의 차이를 줄이기 위해 복원 이미지의 픽셀값과 원본 이미지의 픽셀값의 평균값으로 복원이미지 픽셀값을 결정합니다. 이렇게 계속 학습을 반복하면서 평균값을 구해 그 값으로 대체를 하는 것이죠. 그러면, 최종 복원 이미지는 smoothing 현상으로 인해 그림이 뭉개져서 그려질 것이고, 디테일한 부분을 표현하지 못하게 됩니다.

 

따라서, SRGAN 논문에서는 MSE를 Object Function으로 지정해선 안 된다는 말과 함께 GAN 모델을 사용하며, 새로운 loss function인 Perceptual loss를 제안합니다. 이 loss에서도 MSE가 일부분 사용되긴 하지만, 다른 loss function도 사용되어 smoothing 현상이 많이 완화됩니다. 그래서, CNN 계열 모델들보다 디테일한 부분을 더 잘 표현할 수 있게 됩니다.

 

본론 - 1 (모델 선택의 고충)

이 글에는 SRGAN을 구현하며 느낀 고충과 결과를 남기겠습니다.

 

먼저, 구현 시 참고했던 자료는 이 github 자료입니다. 위 github자료를 통해 SRGAN 모델을 구현한 과정에서 느낀 고충은 Colab 환경에서 GPU 사용량의 한계였습니다.

• 1) SRGAN 모델을 이번 프로젝트에 사용할 데이터셋으로 학습 시 20 에폭 돌리는데 3시간이 넘게 걸리며, GPU 사용량이 끊어짐. 그다음에 10 에폭만 돌려서 학습을 시켰지만, 결과가 너무 좋지 않음.
• 2) Fine-Tuning을 시도했지만, 이 역시 1 에폭에 20분이 넘는 시간이 걸리며 GPU의 사용량이 먼저 끊기고 맘. 이는 GPU에 맞춰져 있는 github 코드를 고쳐 cpu로 돌렸지만, 학습 속도가 너무 느리며 런타임이 끊어짐.

그래서, 제가 생각했던 것은 일단 결과를 살펴보자는 것이었습니다.

 

1번 문제의 경우는 생각보다 큰 문제가 아니었습니다. 아니, 문제 정의 자체가 잘못되었습니다. 제가 하고자 하는 것은 사용자가 자신이 초고해상도로 복원하고자 하는 이미지를 사이트에 업로드하면, 이미지를 초고해상도로 복원해 주며 사용자의 필요에 따라 클래스를 예측해 주는 것이었습니다. 만약, 제가 가지고 있는 bird dataset으로 직접 모델 훈련을 진행한다면, 모델이 일반화되지 못할 것입니다. 즉, 사용자의 입력 사진이 새가 아니라면, 좋은 성능을 보이지 못할 것입니다. 뿐만 아니라, 클래스도 예측할 수 없겠죠. 따라서, Training을 직접 하는 것은 문제 정의 자체가 잘못된 것이었습니다.

 

그러면, 남은 방법은 직접 데이터를 수집하여 수많은 클래스가 있는 데이터셋을 만들거나 사전에 수집되어 있는 대용량 데이터셋을 이용하여 학습을 시키는 것이었습니다. 그러나, 학습 자체도 GPU의 메모리 부족으로 진행되지 않는 상태였죠. 그래서, 대용량 데이터셋으로 사전학습 시킨 모델을 찾아보다가 하나의 github를 찾았지만, 성능이 기대에 미치지 못했습니다. 

출처: https://github.com/aitorzip/PyTorch-SRGAN

 

 

그래서, 일단 Super Resolution에서 많이 사용되는 데이터셋인 DIV2K dataset을 사용하여 사전학습된 모델을 가져와 사용하고, 결과를 살펴봤습니다.

DIV2K Dataset Description

• 학습 데이터: 800개의 고해상도 이미지에서 시작하여 해당 저해상도 이미지를 얻고 2, 3, 4개의 다운스케일링 요소에 대해 고해상도 및 저해상도 이미지를 모두 제공합니다.
• 검증 데이터: 100개의 고해상도 이미지를 사용하여 저해상도 이미지를 생성합니다.
• 테스트 데이터: 100개의 다양한 이미지를 사용하여 낮은 해상도에 해당하는 이미지를 생성합니다.

DIV2K Dataset 예시 3장

 

위처럼, 저희가 학습시키려는 bird에 대한 사진들도 있고, 사람과 강아지 등 여러 클래스가 존재해 데이터셋 자체는 크지 않지만, 모델이 일반화될 수 있겠다는 기대를 했고, 동시에 모델의 출력 결과가 좋아 성능을 기대했습니다. 데이터셋이 작은 것이 걱정이 되었지만, SRCNN 모델도 91장의 데이터로도 충분히 훌륭한 성능을 보였었습니다.

 

본론 - 2 (SRGAN 구현 결과 및 구현의 고충)

github에서 제공된 모델을 그대로 구현만 하면 되었기에 크게 어려움이 없을 거라 생각했습니다. 그런데.. 이게 무슨..?

에러 코드

 

OOM(Out Of Memory) 문제가 발생하여 Inference가 되지 않았습니다. OOM 문제는 GPU의 실행환경인 Cuda의 메모리가 부족하여 생기는 문제인데, 입력 이미지의 크기가 큰 경우에 생기더라고요. 저희가 하고자 하는 데이터는 256 * 256 shape의 이미지였기에 고려하지 않았던 문제인데, 예상치 못한 난관에 부딪혔습니다.

 

따라서, CPU 환경에서 진행하는 것으로 바꾸었습니다. 실행 속도가 걱정이었지만, 어차피 Train을 할 것은 아니기 때문에 코드만 바꿔주면 됐습니다. 사실, github에 올라와있는 대부분의 코드는 cuda를 설정하는 부분이 있습니다. cuda가 avaliabel 한 경우에는 cuda를 사용하고, 그렇지 않으면 CPU를 그대로 사용하도록 설정되어 있기 때문에 GPU가 아니더라도 환경에 맞춰 작동합니다. 하지만, 제가 참고한 코드는 그렇지 않고 GPU에 맞춰져 있는 코드였기에 CPU인 상태에서는 에러가 발생했습니다. 따라서, GPU에 맞춰진 코드를 CPU에 맞춰진 코드로 바꾸는 작업이 필요했습니다.

 

바꾼 부분은 채널이었습니다. GPU에서는 N * C * H * W 형태의 shape을 input으로 가지지만, CPU의 경우는 N * H * W * C 형태의 shape을 input으로 넣어야 합니다. 이에 따라, 사용하는 모델들 전부를 CPU에 맞도록 바꿔주었습니다.

GPU VS CPU

 

이후, 바뀐 코드로 실행한 결과는 아래와 같습니다. 모델의 성능이 생각보다 좋습니다. Input Image의 해상도를 봤을 때, 정말 많이 복원되었습니다. 추가로 주목해야 될 부분은, 제가 서론에서 말씀드린 것처럼, 정량적 평가와 정성적 평가가 다르다는 점입니다. SR분야에서 많이 사용되는 평가지표인 PSNR과 SSIM을 보더라도, 눈으로 보기에는 SRGAN 모델의 결과가 SRCNN, FSRCNN 모델보다 더 좋지만, 평가지표는 반대의 결과를 보여줍니다.

실행 결과 - 정성적 평가

결론

SRGAN 모델을 구현하며 느낀 고충과 결과를 살펴봤습니다. 생각보다 결과가 잘 나왔지만, 사실 아직 완벽한 것은 아닙니다. Fine-tuning까지 진행한다면, 더 좋은 결과를 기대할 수 있겠지만 GPU의 문제로 그렇지 못했습니다. 심지어, Colab Pro를 결제했지만, 다른 프로젝트와 병행하니 컴퓨팅 단위가 부족하더라고요. 아쉽습니다 ,,

 

다음에는 SRGAN을 개선한 ESRGAN과 REAL-ESRGAN의 결과를 보여드리며 최종적으로 선택한 모델을 보여드리고, Super Resolution에 대한 블로깅을 마치겠습니다.

 

긴 글 읽어주셔서 감사합니다.