[Semantic Seg] U-Net, U-Net++, U-Net 3+
U-Net
등장배경
1. 의료 계열 데이터 부족
환자의 개인정보 issue와 데이터가 있어도 일반인이 labeling하기 힘들어서 데이터 수가 부족
- 이로 인하여 일반적인 Deep Learning으로는 턱없는 데이터 수이다
2. Cell Segmentation (인접한 세포 구분) 문제
세포의 경우 경계선이 모호하여서
경계선(테두리)를 판단하기 어려워 일반적인 Semantic Segmentation으로는 작업 불가능
위의 그림과 같은 세포의 경계를 테두리로 만들어야 하는 작업이 필요하다
Architecture
Contracting Path(Encoder)와 Expanding Path(Decoder)가 대칭인 U자 형을 이룬다
1. Contracting Path (Encoder)
- 입력이미지의 전반적인 특징 추출 -> Context를 만들어 준다
- Max-pooling을 이용하여 Down-sampling (차원 축소)
- Architecture : (3 x 3 convolution Network + Batch Normalization + ReLU) 2개씩
=> Batch Normalization은 구현체 마다 넣거나 뺀다.
정의
def CBR2d (in_channels, out_channels, kernel_size =3, stride =1, padding=1, bias =True):
return nn.Sequential(
nn.Conv2d(in_channels=in_channels,
out_channels=out_channels,
kernel_size =kernel_size,
stride =stride,
padding=padding,
bias=bias),
nn.BatchNorm2d(num_features=out_channels), #구현체에 따라 없기도 한다
nn.ReLU()
)
구현
#skip connection을 위해서 결과를 따로 저장한다
self.enc1_1 = CBR2d(in_channels=1,
out_channels=64,
kernel_size=3,
stride=1,
padding=0, #padding을 0을 주어 크기의 변동이 생김
bias=True
)
#skip connection을 위해서 결과를 따로 저장한다
self.enc1_2 = CBR2d(in_channels=1,
out_channels=64,
kernel_size=3,
stride=1,
padding=0, #padding을 0을 주어 크기의 변동이 생김
bias=True
)
self.pool1 = nn.MaxPool2d(kernel_size=2) #이로 인해 크기가 1/2로 줄음
#skip connection을 위해서 결과를 따로 저장한다
self.enc2_1 = CBR2d(in_channels=64,
out_channels=128, #channel을 두배씩 키운다
kernel_size=3,
stride=1,
padding=0, #padding을 0을 주어 크기의 변동이 생김
bias=True
)
#skip connection을 위해서 결과를 따로 저장한다
self.enc2_2 = CBR2d(in_channels=128,
out_channels=128,
kernel_size=3,
stride=1,
padding=0, #padding을 0을 주어 크기의 변동이 생김
bias=True
)
self.pool2 = nn.MaxPool2d(kernel_size=2) #이로 인해 크기가 1/2로 줄음
- 이런 식으로 enc1_2 ~ enc4_2로 따로 저장하는 것 때문에 메모리에 부담이 생길 수도 있다
- channel 수는 1 -> 64 -> 128 -> 256 -> 512 -> 1024 까지 증가한다
2. Crop & Concate (Skip Connection)
- 아래에서 추가적인 설명 (Encoder와 Decoder를 연결해주는 부분)
구현
#Encoder 5 and Decoder 5
self.enc5_1 = CBR2d(in_channels=512,
out_channels=1024, #channel을 두배씩 키운다
kernel_size=3,
stride=1,
padding=0, #padding을 0을 주어 크기의 변동이 생김
bias=True
)
self.enc5_1 = CBR2d(in_channels=1024,
out_channels=1024, #channel을 두배씩 키운다
kernel_size=3,
stride=1,
padding=0, #padding을 0을 주어 크기의 변동이 생김
bias=True
)
self.up_conv4 = nn.ConvTranspose2d(in_channels=1024,
out_channels=512, #channel을 1/2씩 줄인다
kernel_size=2,
stride=2,
padding=0,
bias=True
) #이미지 크기가 두배로 늘어난다
3. Expanding Path (Decoder)
- Localization을 가능하게 한다
- Transposed Convolution을 통해 Upsampling(차원 늘림)을 진행
- Architecture : 2 x 2 Transposed Convolution 사용
=> Feature Map 크기가 2배 증가한다
=> Channel의 수를 절반으로 줄인다
- Contracting Path에서 저장한 Feature Map과 Concat을 진행한다
정의
#이전 out_channels= 512이지만 skip connection에서 넘어오는 enc4_2과 합쳐져 in_channel은 1024로 들어온다
self.dec4_2 = CBR2d(in_channels=1024,
out_channels=512, #channel을 1/2 줄인다
kernel_size=3,
stride=1,
padding=0,
bias=True
)
self.dec4_1 = CBR2d(in_channels=512,
out_channels=512,
kernel_size=3,
stride=1,
padding=0,
bias=True
)
#out_channels= 256이지만 skip connection에서 넘어오는 놈과 합쳐져 전체 out_channel은 512로 넘어간다
self.upconv3 = ConvTranspose2d(in_channels=512,
out_channels=256,
kernel_size=2,
stride=2,
padding=0,
bias=True
)
출력부
정의
#이전 out_channels= 64이지만 skip connection에서 넘어오는 enc1_2과 합쳐져 in_channel은 128로 들어온다
self.dec4_2 = CBR2d(in_channels=128,
out_channels=64, #channel을 1/2 줄인다
kernel_size=3,
stride=1,
padding=0,
bias=True
)
self.dec4_1 = CBR2d(in_channels=64,
out_channels=64,
kernel_size=3,
stride=1,
padding=0,
bias=True
)
# (1x1)conv를 이용하여서 차원 수를 맞춰준다
self.score_fr = nn.Conv2d(in_channels=64,
out_channels=2, #num_classes
kernel_size=1, # 1x1 kernel
stride=1,
padding=0,
bias=True
)
보충설명. Skip Connection (Crop & Concate)
def forward(x):
... (중 략) ...
#Encoder 4
enc4_1 = self.enc4_1(pool3) #Encoder의 maxpooling에서 나온 값이 들어간다
enc4_2 = self.enc4_2(enc4_1)
pool4 = self.pool4(enc5_2)
... (중 략) ...
#Encoder 5, Decoder 5
enc5_1 = self.enc5_1(pool4) #Encoder의 maxpooling에서 나온 값이 들어간다
enc5_2 = self.enc5_2(enc5_1)
upconv4 = self.upconv4(enc5_2)
#concat
crop_enc4_2 = crop_img(enc4_2, upconv4.size()[2]) #resolution 크기가 달라 crop 진행
cat4 = torch.cat([upconv4, crop_enc4_2], dim =1)
저기서 사용된 crop_img가 무엇일까?
def crop_img(in_tensor, out_size):
'''
Args:
in_tensor(tensor) : tensor to be cut
out_size(int) : size of cut
'''
dim1m dim2 = in_tensor.size()[2:]
out_tensor = in_tensor[:,:,
int((dim1-out_size)/2): int((dim1+out_size)/2),
int((dim1-out_size)/2): int((dim1+out_size)/2),
]
return out_tensor
Skip connection에서 넘어오는 Feature: (64 x 64 x 512)
Transposed Convolution에서 넘어오는 Feature: (56 x 56 x 512)
이 둘을 concat을 하고자 하지만 64, 56 이 resolution(해상도)가 달라서 합칠 수가 없음
=> 방법 64 x 64의 외각을 crop하여 56 x 56으로 만들자
{ (64 x 64 x 512) -(crop)-> (56 x 56 x 512) } concat { (56 x 56 x 512) }
= (56 x 56 x 1024)
보충설명)
이렇게 자른다음
추가적인 Techniques
- Data Augmentation : Random Elastic Deformations
- Pixel-wise loss weight 계산을 위한 Weight map 생성
1. Random Elastic Deformations
Elastic : 탄성(물체에 외부적인 압력을 주어 변형) -> ex) 세포의 모양을 다양하게 만들어줌
2. Weight map
- 경계부분을 더 잘 분류하기 위해서 weight map을 사용했다
- 같은 클래스를 가지는 인접한 셀을 분리하기 위해 해당 경계부분에 가중치를 제공
(d1(x) + d2(x) ) -> 인접한 셀과의 거리
w(x) : 인접한 cell일 수록 가중치를 더 줌 (loss를 더 발생시킴)
거리가 멀면 d1, d2가 증가하고 그러면 exp에 마이너스가 있어서 exp는 낮아지고 이로 인해 weight도 작아진다
한계점
1. U-Net은 기본적으로 깊이가 4로 고정된다
-> 데이터셋마다 최적의 깊이를 찾아야 하는 비용이든다
2. 단순한 Skip connection
-> 동일 깊이의 Encoder와 Decoder만 concat 시키는 단순한 방법이다
이 두 문제를 해결하고자 U-Net++가 등장했다
U-Net++
- 다음과 같이 Encoder를 공유하는 다양한 깊이의 U-Net을 생성
- Skip connection을 동일한 깊이에서의 Feature Maps이 모두 결합하도록 유연한 Feature Map 생성
- Depth가 1 ~ 4인 4가지 구조의 U-Net을 가진 구조라 생각이 가능하다
그래서 이걸 어떻게 계산하는데!
Idea : DenseNet에서 사용되었던 Skip Connection 구조를 사용
Dense Skip Connection
다음과 같은 과정을 거친다 [출처 : Naver Connect Foundation]
먼저 용어 정리
과정
이렇게 하면 Depth 1 + Depth 2 + Depth 3 + Depth 4인 모델을 Ensemble하는 효과
(심지어 그냥 4개를 ensemble하는 것 보다 더 좋은 효과를 보인다고 한다)
그래서 이걸 어떻게 이용하는데!
이걸 Deep Supervision에서 사용을 하는데 이를 설명하려면 U-Net++에서 사용하는 Hybrid loss를 먼저 알아야한다
Hybrid Loss
용어 정리
계산 방법
즉, Pixel wise Cross Entropy만 사용하는 것이 아닌 Dice Coefficient도 사용을 한다
Deep Supervision
Loss 계산을 X0.4 뿐만 아니라 X0.1, X0.2, X0.3, X0.4 4가지 모두 loss로 확용
X0.1, X0.2, X0.3, X0.4 를 모두 Hybrid loss를 구한 다음
4개의 loss를 평균 내서 사용한다
이렇게 U-Net의 한계점인
1. U-Net은 기본적으로 깊이가 4로 고정된다
2. 단순한 Skip connection
를 해결하고자 하였다
한계점
- 많은 Skip connection으로 인한 Memory 증가
- Dense한 구조로 인한 Parameter 증가
- 아래 그림과 같이 동일한 계층에 대해서만 Connection을 진행
- Full scale에서 충분한 정보를 탐색하지 못해 위치와 경계를 명시적으로 학습하지 못함
그래서 U-Net 3+ 가 등장
U-Net 3+
U-Net ++ 의 문제점
- 같은 계층의 정보로만 Skip connection을 진행한다
- parameter가 너무 많고 메모리 사용이 너무 많다
U-Net 3+의 구조
1. 같은 계층의 정보로만 Skip Connection을 진행에 대한 문제
=> Full-scale Skip Connections으로 해결
Full-scale Skip Connections의 3가지 요소
1. Conventional
2. Inter
3. Intra
※ 아래의 모든 예시는 X3으로 가정
1. Conventional
- encoder layer로 부터 Same-scale의 feature maps을 받음 (U-Net ++ 와 동일)
- Simple Full-scale Skip Connections
- 크기문제는 따로 발생하지 않음 (같은 계층이므로)
- 채널 수를 맞추기 위해 3x3 Conv 에 Out_channel 64를 사용(추후 이유 설명)
2. Inter
- encoder layer로 부터 Small-scale (receptive field느낌) 의 feature maps을 받음
(여기서 Small scale은 이미지 크기가 작다가 아닌 하나의 픽셀이 담고있는 원본 이미지의 크기가 작다는 것이다)
- 즉, 본인보다 이미지 크기가 큰 부분에서 low level의 feature map을 받음
- Inter Full-scale Skip Connections
- 풍부한 공간 정보를 통해 경계 강조
- 크기문제를 각각 maxpooling을 4 (1/4배 줄임), 2 (1/2배 줄임)를 두어서 해결하였다
- 채널 수를 맞추기 위해 3x3 Conv 에 Out_channel 64를 사용(추후 이유 설명)
- ex) X3의 Decoder의 경우 X1의 Encoder와 X2의 Encoder의 정보가 전달된다
3. Intra
- decoder layer로 부터 (마지막 encoder(depth 5)포함) large-scale (receptive field느낌) 의 feature maps을 받음
(여기서 large scale은 이미지 크기가 크다가 아닌 하나의 픽셀이 담고있는 원본 이미지의 크기가 크다는 것이다)
- 즉, 본인보다 이미지 크기가 작은 부분에서 high level의 feature map을 받음
- Intra Full-scale Skip Connections
- 어디에 위치하는지 위치 정보 구현
- 크기문제를 각각 billinear upsample을 4 (4배 늘림), 2 (2배 늘림)를 두어서 해결하였다
- 채널 수를 맞추기 위해 3x3 Conv 에 Out_channel 64를 사용(추후 이유 설명)
- ex) X3의 Decoder의 경우 X5의 Encoder와 X4의 Decoder의 정보가 전달된다
"이로써 같은 계층의 정보로만 Skip Connection하는 문제를 해결하였다"
2. parameter가 너무 많고 메모리 사용이 너무 많다는 문제
=> 모든 Decoder의 channel수를 320으로 통일 하여서 해결
위에서 언급했듯이 모든 Encoder layer에서 Skip connection과정에서 3x3 Conv의 Out_channel을 64 channel로 고정하여
5개의 channel이 합쳐져서(concat) 320으로 유지되도록 구성하였다
추가적인 테크닉Noise나 False Positive 문제에 대한 해결
Classification-guided Module (CGM)을 추가하였다
다음과 같은 구조를 채택을 하면 low-level layer에서 넘겨지는 Feature map에
Background의 noise가 발생하여서 만은 False Positive문제가 발생한다
(False Positive : 물체가 없는데 물체가 있다고 판단하는 것)
이를 개선하고자 Extra classification task (추가적인 분류 작업)을 진행하였다
Idea : 가장 깊은 층인 (High level feature maps)인 X5 decoder를 활용하여서 물체가 있는지를 판단
방법
1. X5 Decoder에 Dropout - 1x1 Conv - AdaptiveMaxPool - Sigmoid를 통과 시킨다
2. 1의 결과로 Classification의 확률을 만든다 (물체가 있는지 없는지)
3. 2에서 구한 확률을 Argmax을 통해서 물체가 있으면 1 , 없으면 0을 뱉어 낸다
4. 3에서 나온 1 또는 0 을 low-layer마다 곱한다
(물체가 없다면 0이 곱해져서 False Positive가 해결된다)
=> 즉, high level에서 Object가 없다고 판단을 하면 전부 없애는 방법
경계 부분을 잘 학습하기 위한
추가적인 테크닉
3가지추가적인 테크닉 Loss Function을 Summation하여서 사용한다
Hybrid Loss = Focal Loss + ms-ssim Loss + IoU
각각의 loss의 역할
- Focal loss : 클래스의 불균형 해소
- ms-ssim Loss : Boundary 인식 강화
- IoU : 픽셀의 분류 정확도 상승
요약 (4가지 Techniques)
- Full-scale Connections
- Classification-guided Module (CGM)
- Full-scale Deep Supervision
- Hybrid loss Function (Focal + ms-ssim + IoU)
U-Net의 Encoder만 바꾼 모델들 (Residual U-Net, Mobile U-Net, Eff U-Net)
1. Residual U-Net
Encoder 와 Decoder 부분의 block마다 residual unit을 추가
2. Mobile U-Net
좀더 빠른 속도를 위해서 고안되었다
MobileNet을 Backbone으로 사용하여 더 빠르고 가볍게 제작
- 참고자료 : https://medium.com/ddiddu-log/mobiole-unet%EC%9D%84-%ED%99%9C%EC%9A%A9%ED%95%9C-%EC%8B%A4%EC%8B%9C%EA%B0%84-%EB%B0%B0%EA%B2%BD-%EC%A0%9C%EA%B1%B0-%EB%AA%A8%EB%8D%B8-eb066cc29524
- github: https://github.com/anilsathyan7/Portrait-Segmentation
3. Eff U-Net
EfficientNet을 사용
추가적으러 Encoder에는 EfficientNet을 Decoder에는 Skip Connection을 적용하였다
관련 논문
U-Net
https://arxiv.org/pdf/1505.04597.pdf
U-Net ++
https://arxiv.org/pdf/1912.05074.pdf
U-Net 3+
https://arxiv.org/ftp/arxiv/papers/2004/2004.08790.pdf
EfficientUNet
DenseUNet
https://arxiv.org/pdf/1611.09326.pdf
ResidualUNet
https://arxiv.org/pdf/1711.10684.pdf