[논문 리뷰] Crafting a Toolchain for Image Restoration by Deep Reinforcement Learning

Ke Yu, Chao Dong외 2명|arXiv (Cornell University)|2018. 04. 10.

Advanced Image Processing Techniques참고 문헌 41인용 수 31

한 줄 요약

이 논문은 복잡하고 혼합된 왜곡을 가진 이미지를 복원하기 위해 경량의 작업별 특화 컨볼루션 네트워크로 구성된 도구 상자에서 동적으로 선택하는 강화학습 프레임워크인 RL-Restore를 제안한다. 단계별 PSNR 기반 보상과 에이전트 및 도구의 동시 훈련을 통해 순차적 결정 문제로 복원을 공식화함으로써, 대규모 CNN과 유사한 성능을 달성하면서도 계산량을 82.2% 줄였고, 예측되지 않은 아티팩트에 대한 강건성도 향상시켰다.

ABSTRACT

We investigate a novel approach for image restoration by reinforcement learning. Unlike existing studies that mostly train a single large network for a specialized task, we prepare a toolbox consisting of small-scale convolutional networks of different complexities and specialized in different tasks. Our method, RL-Restore, then learns a policy to select appropriate tools from the toolbox to progressively restore the quality of a corrupted image. We formulate a step-wise reward function proportional to how well the image is restored at each step to learn the action policy. We also devise a joint learning scheme to train the agent and tools for better performance in handling uncertainty. In comparison to conventional human-designed networks, RL-Restore is capable of restoring images corrupted with complex and unknown distortions in a more parameter-efficient manner using the dynamically formed toolchain.

연구 동기 및 목표

• 단일 대규모 CNN의 복잡하고 혼합된 이미지 왜곡을 다룰 때 발생하는 높은 파라미터 및 계산 비용의 한계를 해결하기 위해.
• 경량의 작업별 특화 네트워크를 순차적으로 조합함으로써 적응적이고 파라미터 효율적인 복원을 가능하게 하기 위해.
• 에이전트와 도구의 동시 훈련을 통해 중간 복원 상태에서의 불확실성과 아티팩트에 대한 강건성을 향상시키기 위해.
• 가시적인 도구 체인 선택을 통해 단계별 복원 과정을 드러내어 해석 가능성(해석 가능성)을 제공하기 위해.

제안 방법

• 프레임워크는 특정 작업(이면화, 노이즈 제거, JPEG 아티팩트 감소)을 위해 훈련된 소형 전문화된 CNN 도구들로 구성된 도구 상자를 사용한다.
• 에이전트는 현재 이미지 상태와 이전 행동을 바탕으로 단계별 보상 함수(PSNR 향상 비례)에 따라 도구를 순차적으로 선택하는 정책을 학습한다.
• 에이전트와 도구를 동시에 훈련하는 공동 학습 전략을 통해 중간 복원 상태의 불확실성과 아티팩트를 더 잘 다룰 수 있도록 한다.
• 이미지 품질이 만족스럽다고 판단되면 에이전트가 자동으로 복원 과정을 중단하여 과적합을 방지하고 계산량을 절약할 수 있다.
• 보상 함수는 PSNR, SSIM, 최종 PSNR, 최종 MSE를 사용해 평가되며, 단계별 PSNR가 PSNR 및 SSIM 지표에서 가장 우수한 성능을 보였다.
• 도구의 훈련 데이터에는 합성 노이즈와 압축을 포함하여 중간 아티팩트에 대한 강건성을 향상시켰다.

실험 결과

연구 질문

• RQ1 강화학습 에이전트가 복잡하고 혼합된 왜곡을 가진 이미지를 복원하기 위해 경량 네트워크의 동적 도구 체인을 효과적으로 조합할 수 있는가?
• RQ2 에이전트와 도구의 동시 훈련이 알려지지 않은 또는 중간 아티팩트에 대한 강건성을 어떻게 향상시키는가?
• RQ3 다양한 보상 함수가 복원 성능과 수렴에 미치는 영향은 무엇인가?
• RQ4 자동 중단 메커니즘이 복원 품질과 계산 효율성에 어떤 영향을 미치는가?
• RQ5 이 프레임워크는 단일 대규모 CNN보다 예측되지 않은 왜곡에 더 잘 일반화되는가?

주요 결과

• RL-Restore는 FLOPs를 82.2% 줄였음에도 불구하고 단일 대규모 CNN과 유사한 PSNR 성능을 달성하여 높은 계산 효율성을 입증했다.
• 공동 훈련 전략은 기본 도구 대비 0.2 dB 향상된 PSNR를 기록했고, 훈련 데이터에 노이즈 증강을 추가하면 추가로 0.2 dB 향상되었다.
• 단계별 PSNR 보상 함수는 PSNR 및 SSIM 지표에서 단계별 SSIM, 최종 PSNR, 최종 MSE보다 우수한 성능을 보였다.
• 자동 중단 기능을 제거하면 평균적으로 PSNR가 0.15 dB 감소했으며, 경미한 왜곡에서는 더 큰 감소 폭를 보여 과복원 위험을 시사했다.
• 12단계의 도구 체인과 3종의 도구 조합이 최적의 성능을 보였으며, 더 긴 체인은 훈련 복잡도로 인해 성능 향상이 유의미하지 않았다.
• 프레임워크는 해석 가능한 단계별 복원 과정을 드러내어 블랙박스 CNN에 비해 투명성을 제공한다.