[논문리뷰] U-Net 논문 리뷰

U-Net 논문 리뷰입니다.

본 논문은 2015년도에 등장한 논문입니다.

논문 : U-Net: Convolutional Networks for Biomedical Image Segmentation

저자 : Olaf Ronneberger, Philipp Fischer, and Thomas Brox

[1] Background

Deep CNN Network는 CV Task에서 뛰어난 성능을 보여왔습니다.

 

ImageNet 와 같은 규모 데이터셋의 등장과 함께 딥러닝의 성능이 크게 향상되었습니다.

 

이러한 CNN 은 보통 이미지 분류 작업에 쓰이는데, 의료 도메인에서는 이미지 분류만으로는 충분하지 않으며 이미지의 각 픽셀이 어떤 상태인지, 어떤 구조인지를 알아야 하는 경우가 많아 적용되기 힘든 문제가 있습니다. (이미지 내에 여러 세포가 들어있기 때문에 픽셀별로 클래스 분류를 해야하는 거죠.)

 

또한, 의료 데이터는 확보하기 어려운 경우가 많아 대량의 학습 데이터를 사용하는 것이 쉽지 않습니다.

 

즉, 의료 도메인에서 적은 양의 이미지 데이터로도 학습과 예측이 용이한 모델을 만드는 것이 필요한 상황이었습니다.

 

Patch

 

기존 연구들에서는 영역에 대해 학습할 때 Sliding Window 방식을 사용한 경우가 많았습니다.

 

하지만 이 Sliding Window 방식에는 두 가지 단점이 있는데,

 

1) Conv 연산이 각 영역마다 실행되어야 하고, 중복되는 영역이 있어서 비효율적이며 속도가 느립니다

2) 위치 정보의 정확성과 문맥 사용 간에 트레이드오프가 존재합니다

→ 큰 ROI는 전체적으로 보기 때문에 세부 정보를 놓치고 작은 ROI는 세부적으로 보기 때문에 주변 문맥 정보를 놓친다고 이해

ex) "이 픽셀은 단순히 검은색이다"로 판단하지만, 이 픽셀이 고양이의 귀라는 사실은 놓침. (작은 ROI)

ex) "전체적으로 고양이처럼 보인다"라고 판단하지만, 고양이 귀의 정확한 위치는 흐릿하게 됨. (큰 ROI)

 

따라서, 이를 보완하고자 Sliding Window 방식이 아닌 Fetch 방식을 사용한 모델 U Net이 개발되었습니다.

 

---

[2] Architecture

U-Net Architecture

 

위 사진이 U-Net의 기본 구조이고 수축 레이어(왼쪽)와 확장 레이어(오른쪽)로 구성되어 있습니다.

 

기본적으로 수축 레이어에서는 데이터를 압축해서 특징을 추출하고, 확장 레이어에서는 세분화된 결과를 만듭니다.

 

Contradiction Path

Contradiction Path

 

수축 레이어부터 좀 더 자세히 보면 FC Layer를 제외한 전형적인 Convolution Network를 수행합니다.

그래서 각 부분 마다 2개의  Conv 연산, BN, ReLU,  Max Pooling (Stride 2, 2x2) 다운샘플링 연산을 반복합니다.

 

다운샘플링 단계마다 이미지의 크기는 줄이지만 특징 정보의 표현을 더 풍부하게 하기 위해 채널 수를 늘립니다.

 

그리고 Expansive Path 로 가기 전에 1024 채널로 출력되는 Conv Layer가 두 번 사용되는데 그 이유는

Conv Layer를 두 번 사용함으로써 데이터의 추상적인 특징을 더 잘 학습하게 하고 비선형성을 추가할 수 있다고 합니다.

 

Expansive Path

Expansive Path

 

 Expansive Path의 각 단계는 채널 수를 절반으로 줄이는 UP Conv Layer, 두 번의 3x3 Conv Layer, Skip Connection으로 구성됩니다.

 

Skip Connection이란, Contradiction Path에서 추출된 Feature Map을 업샘플링된 벡터에 Concat시켜서 연산하는 부분입니다.

 

이렇게 Feature Map을 가져올 때 전체 Feature Map을 가져오는 게 아니라 가운데 일정 부분을 Crop해서 가져오는데

이는, Conv 연산에서 경계 픽셀이 손실되기 때문에 필요하다고 합니다.

💡Padding 없이 Conv 연산을 수행하기 때문에 가장자리 엣지 부분의 정보가 손실될 수 있습니다.
반대로 비교적 중앙 부분은 모든 Conv Filter가 완전하게 적용된 영역이기 때문에 정보 손실이 적어 신뢰도가 높죠.
그래서 신뢰도가 높은 중앙 영역을 Crop 해서 Expansive Path의 같은 층에 전달한다고 이해했습니다.

 

그렇게 Upsampling과 두 번의 Conv 연산, Skip Connection을 거쳐 마지막 레이어에는 1x1 conv Layer로 Feature Map을 원하는 클래스 수에 맞게 분류 / Segmentation 합니다.

 

U-Net에는 총 23 개의 Conv Layer가 포함되어 있어 깊이가 깊기 때문에 복잡한 데이터를 처리할 수 있습니다.

 

Overlap tile Mirroring

 

여기서 Overlap tile Mirroring 이란,  큰 이미지를 작은 타일로 나눠 처리할 때, 경계 영역의 문맥 정보를 보완하기 위해 사용되는 기법을 의미합니다.

 

위 사진에서 타일의 경계 부분인 노란색 영역은 전체 문맥 정보를 제대로 반영하지 못할 수 있습니다. (No padding)

 

그래서 타일 경계 외부의 문맥 정보를 알기 위해 이미지를 반사(mirroring) 시켜 경계 외부의 데이터를 생성합니다.이 반사된 데이터는 경계 바깥의 실제 정보를 흉내내어 네트워크가 더 정확히 경계 부분을 처리할 수 있도록 돕습니다.

 

이러한 미러링 기법을 적용하면 경계 부분에서 끊김 없는 매끄러운 결과를 얻을 수 있습니다.→ 당시에 GPU 메모리 제약으로 인해 이미지를 작은 타일로 쪼개야 하지만 타일 경계의 문맥 부족 문제 해결을 위해서는 미러링이 필수적이라고 합니다.

 

---

[3] Training & Data Augmentation

 

a : raw 이미지 (입력된 세포 이미지)

b : label 값 (세분화된 결과)

c : label 이미지를 흑백으로 변환

d : point-wise 변환 값

 

위 이미지의 a와 같은 입력 이미지와 b와 같은 라벨값을 활용해 네트워크를 학습시키며 Caffe 프레임워크의 SGD (확률적 경사하강법) 알고리즘을 활용한다고 합니다.

 

그리고 GPU 메모리를 최대한 활용하기 위해 Batch Size를 줄이고 더 큰 Patch 를 활용하는데,

이럴 경우에 Batch Size 가 작아 최적화가 잘 안되는 상황을 대비해 모멘텀 값을 0.99로 설정해 과거 샘플들이 현재 최적화 단계에 큰 영향을 줄 수 있도록 조치하였습니다.

 

수식으로 알아보기

 

P(x)는 각 픽셀이 특정 클래스에 속할 확률을 계산하며 k는 해당 클래스, x는 픽셀 위치를 나타내며 K는 총 클래스 수를 의미합니다.

 

따라서, 위와 같은 소프트맥스 함수로 가장 높은 수치의 클래스 확률을 최대화하도록 설계되었고 이렇게 각 픽셀마다의 클래스 확률을 계산한 후 예측값과 실제값의 차이를 Cross-Entropy Loss 로 측정합니다.

 

 

다음 수식은 가중치 맵 수식 w(x) 입니다.

 

x 위치 픽셀에 어느정도의 가중치를 부여할 것이냐를 결정하는 수식이구요.

 

Wc(x)는 클래스 빈도를 균형 있게 만드는 가중치이고 d1(x)는 가장 가까운 세포까지의 거리, d2(x)는 두 번째로 가까운 세포까지의 거리를 의미합니다.

 

따라서, 경계 근처 픽셀에 더 높은 가중치를 부여하기 위해 각 픽셀의 세포 경계까지의 거리를 사용합니다.

 

Data Augmentation

 

의료 데이터는 도메인 특성상 방대한 데이터를 구하기가 힘듭니다.

 

이렇게 학습 데이터가 부족한 상황에서는 데이터 증강을 통해 모델의 강건함을 높이는 것이 중요하다고 논문에서 언급하고 있습니다.

 

또한, 현미경 이미지의 경우 회전, 이동, 변형, 흑백 등 다양한 변화가 일어날 수 있어 이러한 변화에 대해 강인함을 길러야 합니다.

 

이를 위해서 논문에서는 Elastic Deformation (탄성 변환) 기법을 활용합니다.

 

위 사진처럼 이미지에 유연한 왜곡을 추가하여 네트워크가 다양한 모양에 적응할 수 있도록 하는 기법이고

이미지를 픽셀 단위로 나눠 각 픽셀의 중심점을 상하좌우 임의로 이동시켜서 왜곡을 준다고 합니다.

 

이동 값은 표준 편차가 10 픽셀인 가우시안 분포에서 샘플링되고, bicubic interpolation (삼차 보간법)으로 계산됩니다.

 

마지막으로, Dropout을 통해 과적합 방지까지 해 주었습니다.

 

---

[4] Experiments

• 현미경으로 찍은 세포를 Segmentation하는 대회
• U-Net은 warping error에서 0.000353으로 최고 성능을 기록하며 1위를 차지

 

• "PhC-U373" 데이터셋에서 U-Net은 92.0%의 IOU를 기록하며, 2위 알고리즘(83%)보다 크게 앞섬
• "DIC-HeLa" 데이터셋에서도 77.5%의 IOU로, 2위 알고리즘(46%)보다 월등히 높은 성능을 보임

 

• Table 2의 추가설명
• a와 b를 비교했을 때 청록색 부분(U-Net의 예측)이 노란색 경계(라벨값)와 매우 일치 (PhC-U373)
• c와 d에서도 무작위 색상으로 표시된 세포 경계가 라벨과 높은 일치도를 보임 (DIC-HeLa)

 

---

[5] Conclusions

• U-Net은 적은 수의 라벨 데이터로도 의료 도메인의 Segmentation 작업에서 뛰어난 성능을 발휘하도록 설계됨
• 마치 U 처럼 생긴 모양의 모델 아키텍쳐를 구성하고 있음
  • Contradiction Path와 Expansive Path로 구성되어 있으며 Overlap Tile Mirroring과 elastic deformation 기법을 적용
• 다양한 데이터셋에서 SOTA 달성

 

---

[6] References

https://www.youtube.com/watch?v=O_7mR4H9WLk&t=635s

https://velog.io/@minkyu4506/%EB%85%BC%EB%AC%B8-%EB%A6%AC%EB%B7%B0-U-Net-%EB%A6%AC%EB%B7%B0

https://pasus.tistory.com/204

 

[논문리뷰] Faster R-CNN 논문 리뷰