[논문 리뷰] ISTA-Net++: Flexible Deep Unfolding Network for Compressive Sensing

Di You, Jingfen Xie외 1명|arXiv (Cornell University)|2021. 03. 22.

Sparse and Compressive Sensing Techniques참고 문헌 26인용 수 8

한 줄 요약

ISTA-Net++는 조건 모듈을 통해 다양한 샘플링 비율에 동적으로 적응하는 유연한 딥 언폴딩 네트워크를 압축 감지에 제안하며, 크로스 블록 전략과 균형 잡힌 훈련 데이터셋(Train400)을 활용해 블로킹 아티팩트를 감소시키고 다양한 이미지 장면에서의 강건성을 향상시킨다. 이는 다수의 데이터셋에서 최신 기술 수준의 성능을 달성하며, 다중 비율 및 다중 장면 작업에서 일관된 성능 향상을 보인다.

ABSTRACT

While deep neural networks have achieved impressive success in image compressive sensing (CS), most of them lack flexibility when dealing with multi-ratio tasks and multi-scene images in practical applications. To tackle these challenges, we propose a novel end-to-end flexible ISTA-unfolding deep network, dubbed ISTA-Net++, with superior performance and strong flexibility. Specifically, by developing a dynamic unfolding strategy, our model enjoys the adaptability of handling CS problems with different ratios, i.e., multi-ratio tasks, through a single model. A cross-block strategy is further utilized to reduce blocking artifacts and enhance the CS recovery quality. Furthermore, we adopt a balanced dataset for training, which brings more robustness when reconstructing images of multiple scenes. Extensive experiments on four datasets show that ISTA-Net++ achieves state-of-the-art results in terms of both quantitative metrics and visual quality. Considering its flexibility, effectiveness and practicability, our model is expected to serve as a suitable baseline in future CS research. The source code is available on https://github.com/jianzhangcs/ISTA-Netpp.

연구 동기 및 목표

• 실제 응용에서 다수의 샘플링 비율을 처리할 때 기존 딥 러닝 기반의 압축 감지 방법의 유연성 부족 문제를 해결한다.
• 훈련 데이터의 데이터셋 불균형을 완화시켜 다양한 이미지 장면, 특히 '희귀로 간주되는' 장면에서의 재구성 강건성을 향상시킨다.
• 새로운 크로스 블록 전략을 통해 재구성된 이미지의 블로킹 아티팩트를 감소시킨다.
• 재훈련 없이도 다중 비율 압축 감지 재구성을 수행할 수 있는 통합 모델을 개발하여 실용성과 효율성을 향상시킨다.

제안 방법

• 샘플링 비율 정보를 각 네트워크 단계에 주입하기 위해 조건 모듈을 사용하는 동적 언폴딩 전략(DUS)을 도입하여 단일 모델이 다수의 CS 비율에 적응할 수 있도록 한다.
• 샘플링 비율에 조건이 되는 동적 그래디언트 강하 모듈(DGDM)과 동적 프락시 매핑 모듈(DPMM)을 활용하여 재구성 과정 중 적응형 최적화를 가능하게 한다.
• 접한 이미지 블록 간 기능 공유를 가능하게 하는 크로스 블록 전략(CBS)을 적용하여 블로킹 아티팩트를 억제하고 텍스처 정밀도를 향상시킨다.
• 표준 데이터셋인 Train91에서 흔히 볼 수 있는 주로 '식물' 카테고리 외의 다양한 장면을 포함한 균형 잡힌 데이터셋(Train400)을 사용해 네트워크를 훈련시킨다.
• 세 부분으로 구성된 프레임워크를 사용한다: 샘플링(컨볼루션 레이어로 모델링), 초기화(차원 불일치 문제 해결), 재구성(동적 언폴딩과 DUS, CBS를 통한 재구성).
• 모든 구성 요소의 공동 최적화를 가능하게 하기 위해 재구성 오차를 최소화하는 손실 함수를 사용하는 엔드 투 엔드 훈련을 활용한다.

실험 결과

연구 질문

• RQ1 재훈련 없이 단일 딥 러닝 모델이 다양한 샘플링 비율에서 압축 감지 재구성을 효과적으로 처리할 수 있는가?
• RQ2 데이터셋 불균형은 새로운 이미지 장면에서 CS 모델의 일반화에 어떤 영향을 미치며, 이를 완화할 수 있는가?
• RQ3 크로스 블록 전략이 딥 언폴딩 기반의 CS 재구성에서 블로킹 아티팩트를 어느 정도 감소시키는가?
• RQ4 Train400처럼 더 균형 잡힌 훈련 데이터셋을 사용하면 다양한 장면에서 재구성 강건성에 측정 가능한 향상이 이루어지는가?

주요 결과

• Set11에서 50% 샘플링 비율일 때 ISTA-Net++는 평균 PSNR 38.73 dB를 기록하며, 조건 모듈이 없는 베이스라인 대비 10%, 30%, 50% 비율에서 평균 0.26 dB 향상된 성능을 보였다.
• 크로스 블록 전략(CBS)은 조건 모듈이 없는 베이스라인 대비 블로킹 아티팩트를 감소시키고 PSNR를 최대 0.64 dB 향상시켰다.
• 균형 잡힌 Train400 데이터셋으로 훈련된 모델은 30% 비율에서 34.81 dB의 PSNR를 기록했으며, 동일한 에포크 수로 훈련된 비균형 Train91 데이터셋으로 훈련된 모델(34.25 dB)보다 뛰어난 성능을 보였다.
• DIV2K 검증 세트에서 30% 비율일 때, Train400로 훈련된 ISTA-Net++ 모델은 '주택' 장면에서 Train91로 훈련된 모델보다 0.97 dB 높은 성능을 기록하여 희귀 장면에 대한 강건성이 향상됨을 입증했다.
• 동적 언폴딩과 크로스 블록 전략의 조합은 모든 테스트 비율 및 데이터셋에서 일관된 성능 향상을 이끌어내어 양 구성 요소의 효과성을 확인했다.
• 실험 결과, 데이터셋의 균형이 데이터셋 크기보다 더 큰 영향을 미치며, Train400에 더 많은 데이터를 추가(예: Train400+Train91)해도 추가적인 성능 향상이 거의 없음을 확인했다.