수리링의 데이터교실

https://smartest-suri.tistory.com/49

지난 번 포스팅에서 리뷰한 U-Net 논문을 파이토치를 이용한 코드로 구현한 과정을 정리해 보겠습니다.

1. [연습] Class 없이 한줄씩 구현

직관적인 이해를 위해서 파이토치 코드로 클래스 없이 한줄씩 유넷 구조를 구현해 보도록 하겠습니다. 

# 먼저 필요한 모듈을 임포트 해줍니다.

import torch
import torch.nn as nn
import torchvision.transforms.functional as TF

1-1. Contracting path - 인코딩 파트

• 논문에서는 valid padding(패딩 없음)을 사용하지만 코드상의 편의를 위해서 모든 콘볼루션 레이어를 same padding(패딩 1)로 구현하겠습니다.
• 이렇게 할 경우 나중에 skip-connection 파트에서 concatenate할 때 크기가 딱 맞아서 crop할 필요가 없습니다.

input_channels = 3 # 일단 RGB라고 생각하고 3으로 설정하겠습니다.

# 첫번째 블락
conv1 = nn.Conv2d(input_channels, 64, kernel_size = 3, padding = 1)
conv2 = nn.Conv2d(64, 64, kernel_size = 3, padding = 1)
pool1 = nn.MaxPool2d(kernel_size = 2, stride = 2)

# 두번째 블락
conv3 = nn.Conv2d(64, 128, kernel_size = 3, padding = 1)
conv4 = nn.Conv2d(128, 128, kernel_size = 3, padding = 1)
pool2 = nn.MaxPool2d(kernel_size = 2, stride = 2)

# 세번째 블락
conv5 = nn.Conv2d(128, 256, kernel_size = 3, padding = 1)
conv6 = nn.Conv2d(256, 256, kernel_size = 3, padding = 1)
pool3 = nn.MaxPool2d(kernel_size = 2, stride = 2)

# 네번째 블락
conv7 = nn.Conv2d(256, 512, kernel_size = 3, padding = 1)
conv8 = nn.Conv2d(512, 512, kernel_size = 3, padding = 1)
pool4 = nn.MaxPool2d(kernel_size = 2, stride = 2)

1-2. Bottleneck - 연결 파트

• 연결 구간에 해당하는 Bottleneck 파트를 작성하겠습니다.
• 여기까지 오면 최종 채널의 수는 1024가 됩니다.

conv9 = nn.Conv2d(512, 1024, kernel_size = 3, padding = 1)
conv10 = nn.Conv2d(1024, 1024, kernel_size = 3, padding = 1)

1-3. Expanding path - 디코딩 파트

• 디코딩 파트에는 Skip-connection을 통한 사이즈와 필터의 변화에 주목해서 보시면 좋습니다.

# 첫 번째 블락
up1 = nn.ConvTranspose2d(1024, 512, kernel_size = 2, stride = 2)
# 위에 아웃풋은 512지만 나중에 코드에서 skip-connection을 통해 다시 인풋이 1024가 됨
conv11 = nn.Conv2d(1024, 512, kernel_size = 3, padding = 1)
conv12 = nn.Conv2d(512, 512, kernel_size = 3, padding = 1)

# 두 번째 블락
up2 = nn.ConvTranspose2d(512, 256, kernel_size = 2, stride = 2)
conv13 = nn.Conv2d(512, 256, kernel_size = 3, padding = 1) # Skip-connection 포함
conv14 = nn.Conv2d(256, 256, kernel_size = 3, padding = 1)

# 세 번째 블락
up3 = nn.ConvTranspose2d(256, 128, kernel_size = 2, stride = 2)
conv15 = nn.Conv2d(256, 128, kernel_size = 3, padding = 1) # Skip-connection 포함
conv16 = nn.Conv2d(128, 128, kernel_size = 3, padding = 1)

# 네 번째 블락
up4 = nn.ConvTranspose2d(128, 64, kernel_size = 2, stride = 2)
conv17 = nn.Conv2d(128, 64, kernel_size = 3, padding = 1) # Skip-connection 포함
conv18 = nn.Conv2d(64, 64, kernel_size = 3, padding = 1)

# 마지막 아웃풋에서는 1x1 사이즈 콘볼루션을 사용한다고 이야기함.
output = nn.Conv2d(64, 2, kernel_size = 1, padding = 1)

1-4. Forward-pass

포워드 학습 과정을 한줄씩 구현해 보겠습니다.

def unet_forward(x):
    # 인코더 파트
    x = TF.relu(conv1(x))
    x1 = TF.relu(conv2d(x))
    x = pool1(x1)

    x = TF.relu(conv3(x))
    x2 = TF.relu(conv4(x))
    x = pool2(x2)

    x = TF.relu(conv5(x))
    x3 = TF.relu(conv6(x))
    x = pool3(x3)

    x = TF.relu(conv7(x))
    x4 = TF.relu(conv8(x))
    x = pool4(x4)

    # 연결 bottleneck 파트
    x = TF.relu(conv9(x))
    x = TF.relu(conv10(x))

    # 디코더 파트
    x5 = up1(x)
    x = torch.cat([x5, x4], dim = 1) # skip-connection
    x = TF.relu(conv11(x))
    x = TF.relu(conv12(x))

    x6 = up2(x)
    x = torch.cat([x6, x3], dim = 1) # skip-connection
    x = TF.relu(conv13(x))
    x = TF.relu(conv14(x))

    x7 = up3(x)
    x = torch.cat([x7, x2], dim = 1) # skip-connection
    x = TF.relu(conv15(x))
    x = TF.relu(conv16(x))

    x8 = up4(x)
    x = torch.cat([x8, x1], dim = 1) # skip-connection
    x = TF.relu(conv17(x))
    x = TF.relu(conv18(x))

    # 아웃풋 파트
    output = output(x)

    return output

지금까지 클래스 없이 파이토치로 코드를 짜면서 유넷 구조를 직관적으로 이해해 보았습니다.

2. [실전] Class 이용해서 구현

이번엔 파이토치를 사용하는 만큼 클래스를 이용해서 실전 유넷 코드 구현을 해보겠습니다. 먼저 코드를 짜는 과정은 유튜브 [PyTorch Image Segmentation Tutorial with U-NET: everything from scratch baby (Aladdin Persson)]를 참고하여 작성했음을 미리 밝히겠습니다. 

https://youtu.be/IHq1t7NxS8k?si=776huHRjVsIlf_rS

2-1.  모든 블락의 2개 콘볼루션(Conv2d) 레이어를 정의하는 클래스

먼저 unet.py 파일을 만들고 클래스들을 작성하여 파일을 모듈화 하겠습니다.

class DoubleConv(mm.Module):
    def __init__(self, in_channels, out_channels):
        super(DoubleConv, self).__init__()
        self.conv = nn.Sequential(
            nn.Conv2d(in_channels, out_channels, kernel_size = 3, stride = 1, padding = 1, bias = False),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace = True),
            nn.Conv2d(out_channels, out_channels, 3, 1, 1, bias = False),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace = True)
        )

    def forward(self, x):
        return self.conv(x)

• 1장과 마찬가지로 편의를 위해서 valid padding 대신 same padding을 사용합니다. Carvana 대회 1등한 팀도 same padding을 사용했다고 하네요. (저는 다른걸로 할거지만, 영상에서 Carvana 데이터셋을 이용해서 학습을 합니다)
• Batch Normalization이 추가되었습니다. 논문에서 따로 Batch Normalization을 해준다는 언급은 없었는데요. 영상에 따르면 UNet이 2015년도 발표되었고, BatchNorm은 2016년에 고안된 아이디어라서 그렇다고 합니다. 찾아보니 유넷을 구현하는 많은 코드가 Batch Normalization을 추가하여 Gradient Vanishing/Exploding 문제를 보완하는 것 같습니다.
• bias를 False로 설정해 주는 이유는 중간에 BatchNorm2d를 추가해주기 때문입니다. bias가 있어봤자 BatchNorm에 의해서 상쇄(cancel)되기 때문에 굳이 bias가 필요 없다고 영상에서 말하고 있습니다. 이부분은 나중에 BatchNorm 논문을 통해 따로 확인해 보겠습니다.

2-2. 유넷 전체 구조를 정의하는 클래스

해당 파트는 코드가 길어지기 때문에 주석을 이용해서 각 부분을 설명했습니다 :) 비교적 간단하지만 그래도 코드를 이해하시려면 유넷 전체 구조에 대한 이해가 필수적입니다.

• 논문에서는 마지막 아웃풋의 채널이 2였습니다만, 저는 레이블이 0, 1로 binary인 데이터를 다룰 예정이기 때문에 output_channels의 default 값을 1로 두었습니다.
• features 리스트는 각 블록에서 피처맵의 사이즈를 순서대로 정의하는 역할을 합니다. (contracting path의 경우 순서대로 64, 128, 256, 512 - expanding path의 경우 그 반대)

class UNET(nn.Module):
    def __init__(
            self, in_channels = 3, out_channels = 1, features = [64, 128, 256, 512]
    ):
        super(UNET, self).__init__()
        self.downs = nn.ModuleList() # Contracting path - 인코딩 파트의 모듈을 담을 리스트 선언
        self.ups = nn.ModuleList()   # Expanding path - 디코딩 파트의 모듈을 담을 리스트 선언
        self.pool = nn.MaxPool2d(kernel_size = 2, stride = 2) # 풀링은 모든 블럭에서 공통 사용됨

        # Contracting path (Down - 인코딩 파트) ------------------------
        for feature in features:
            self.downs.append(DoubleConv(in_channels, feature)) # 블록마다 더블콘볼루션 해주고 아웃풋은 feature맵 리스트 순서대로 할당(64, 128...)
            in_channels = feature # 다음 모듈의 인풋 사이즈를 feature로 업데이트

        # Bottleneck (인코딩, 디코딩 연결 파트) ---------------------------
        size = features[-1] # 512
        self.bottleneck = DoubleConv(size, size * 2) # 인풋 512 아웃풋 1024

        # Expanding path (Up - 디코딩 파트) ----------------------------
        for feature in features[::-1]: # 피처맵 사이즈 반대로!
            # 먼저 초록색 화살표에 해당하는 up-conv 레이어를 먼저 추가해 줍니다.
            self.ups.append(
                nn.ConvTranspose2d(
                    feature*2, feature, kernel_size = 2, stride = 2
                    # 인풋에 *2 해주는 이유 : 나중에 skip-connection을 통해서 들어오는 인풋 사이즈가 더블이 되기 때문!
                    # kernel과 stride size가 2인 이유는.. 논문에서 그렇게 하겠다고 했음 '_^ 
                )
            )
            # 이제 더블 콘볼루션 레이어 추가
            self.ups.append(DoubleConv(feature * 2, feature))

        # Output (아웃풋 파트) -----------------------------------------
        last_input = features[0] # 64
        self.final_conv = nn.Conv2d(last_input, out_channels, kernel_size = 1)

    ####################### **-- forward pass --** #######################
    def forward(self, x):
        skip_connections = []

        # Contracting path (Down - 인코딩 파트) ------------------------
        for down in self.downs: # 인코딩 파트를 지나면서 각 블록에서 저장된 마지막 모듈 하나씩 이터레이션
            x = down(x)
            skip_connections.append(x) # skip_connection 리스트에 추가
            x = self.pool(x)

        # Bottleneck (인코딩, 디코딩 연결 파트) ---------------------------
        x = self.bottleneck(x)
        skip_connections = skip_connections[::-1] # 디코딩 파트에서 순서대로 하나씩 뽑기 편하게 리스트 순서 반대로 뒤집어주기

        # Expanding path (Up - 디코딩 파트) ----------------------------
        for i in range(len(skip_connections)):
            x = self.ups[i * 2](x) # self_ups에는 순서대로 ConvTranspose2d와 DoubleConv가 들어가 있음. 
            # 0, 2, 4... 짝수번째에 해당하는 인덱스만 지정하면 순서대로 ConvTranspose2d와(up-conv) 모듈만 지정하게 됨
            skip_connection = skip_connections[i] # skip_connection 순서대로 하나씩 뽑아서
            # concatenate 해서 붙여줄(connection) 차례!
            # 그런데 만약 붙일때 shape이 맞지 않는다면... (특히 이미지의 input_shape이 홀수인 경우 이런 뻑이 나게됨)
            if x.shape != skip_connection.shape:
                x = TF.resize(x, size = skip_connection.shape[2:]) # 간단히resize로 맞춰주겠음
            concat_skip = torch.cat((skip_connection, x), dim = 1) # 이제 붙임!
            x = self.ups[i * 2 + 1](concat_skip)
            # 1, 3, 5... 홀수번째에 해당하는 인덱스만 지정하면 순서대로 DoubleConv 모듈만 지정하게 됨
            
        return self.final_conv(x)

2-3. 테스트

• 간단한 랜덤 텐서를 생성해서 우리가 구현한 유넷 모델이 제대로 작동하는지 확인하겠습니다.
• 모델 인풋과 아웃풋의 shape이 정확히 같은지 확인하는 작업이 들어가는데, image segmentation 작업의 특성상 output mask의 크기가 인풋과 동일해야 하기 때문에 그렇습니다.
• assert 문법을 사용해서 인풋 아웃풋 shape이 다른 경우가 감지되면 AsserstionError를 발생시켜서 Debug에 활용하도록 합니다.

# 유넷 모델 테스트하는 함수 작성

def test():
    # 3장의 1채널(grayscale), width height가 각각 160인 랜덤 텐서 생성 (테스트용으로!)
    x = torch.randn((3, 1, 160, 160))
    # in_channels 1로 설정 (그레이스케일 이미지), out_channels 1로 설정 (binary output)
    model = UNET(in_channels = 1, out_channels = 1)
    # forward pass
    preds = model(x)
    # preds와 x의 shape 확인하기 - 두개가 같아야 함
    print(preds.shape)
    print(x.shape)
    assert preds.shape == x.shape 
        # True : 오케이
        # False : AssertionError

if __name__ == "__main__": # 메인 파일에서만 작동하고 모듈로 import된 경우에는 작동하지 않도록 함
    test()
  
# 실행 결과 ---------------------------
# torch.Size([3, 1, 160, 160])
# torch.Size([3, 1, 160, 160])

구동 결과 shape이 정확히 같아서 test 함수가 제대로 작동한 것을 확인했습니다.

여기까지 작성한 코드를 unet.py 파일로 저장해서 모듈화 해주었습니다.

---

이후 데이터를 로드에 필요한 dataset.py를 작성한 다음 학습에 필요한 utils.py train.py를 차례대로 작성하는 순서로 학습을 진행합니다. 해당 파트는 본 포스팅에서 제외하도록 하겠습니다.

지금까지 유넷 논문에서 살펴본 구조를 파이토치 코드로 구현하는 과정을 포스팅해 보았습니다 :-) 다음 포스팅에서는 트랜스포머 코드화 작업을 수행해 보도록 하겠습니다. 감사합니다!

[주의] 본 포스팅은 수리링이 직접 U-Net 논문 원문을 읽고 리뷰한 내용을 담았으며, 참고 문헌이 있는 경우 출처를 명시하였습니다. 본문 내용에 틀린 부분이 있다면 댓글로 말씀해 주시고, 포스팅을 출처 없이 불법 공유하지 말아주시기 바랍니다. 감사합니다.

U-Net: Convolutional Networks for BiomedicalImage Segmentation

https://arxiv.org/pdf/1505.04597

이미지 세그멘테이션(image segmentation)에서 빼놓을 수 없는 근본 모델 유넷. 유넷은 픽셀 기반으로 이미지를 분할하여 구분하는 모델로, 많은 최신 모델 속에서 그 구조가 사용되며 활약하고 있습니다.

유넷은 ISBI cell tracking challenge 2015 대회에서 등장한 모델로, 세포 이미지로부터 뛰어난 세그멘테이션 성능을 보여주며 우승을 차지했습니다. 놀라운 점은, 이 대회에서 제공된 학습 데이터가 고작 이미지 30장이었다고 합니다. 저는 논문 리뷰를 마치고서야 이 사실을 알게 되어서, 논문에서 계속 데이터 수가 부족했다고 우는 소리를 하는 이유가 있었구나.....! 대박적이다...! 라는 생각을 했습니다.

따라서 유넷 논문을 처음 공부하시는 분들께서는, 유넷이 어떻게 이러한 한계를 극복하며 뛰어난 segmentation을 구현할 수 있었는지에 주목하면서 리뷰를 해보시면 좋을 것 같습니다.

Abstract

In this paper, we present a network and training strategy that relies on the strong use of data augmentation to use the available annotated samples more efficiently.

초록에서는 주어진 데이터를 더욱 효율적으로 사용하기 위해 data augmentation - 데이터 증강을 통한 네트워크와 학습 전략을 제시하겠다고 밝히고 있습니다. 

The architecture consists of a contracting path to capture context and a symmetric expanding path that enables precise localization. 

U넷의 네트워크 구조는 크게 2가지 path로 구성됩니다.

1. 수축 경로 'contracting path' : capture context - 넓은 범위에서 이미지의 주요한 특징과 의미 추출
2. 확장 경로 'expanding path' : 정교한 localization - 앞에서 추출한 의미를 위치정보와 결합하여 각 픽셀마다 어떤 객체에 속하는지를 구분
3. (+) 'bottleneck' 연결 구간 : 수축 경로와 확장 경로를 연결 (논문에 언급되어 있지 않지만 흔히 보틀넥이라고 칭합니다.)

We show that such a network can be trained end-to-end from very few images and outperforms the prior best method (a sliding-window convolutional network) on the ISBI challenge for segmentation of neuronal structures in electron microscopic stacks.

구조를 그렇게 설계했더니 아주 적은 수의 이미지로도 end-to-end 학습이 가능했을 뿐더러 기존에 sliding-window 형식의 CNN구조를 사용하고 있던 베스트 모델보다도 성능이 훨씬 뛰어났다는 이야기를 하고 있습니다.

end-to-end 딥러닝이 뭔지 궁금하신 분께서는 아래 유튜브 동영상을 참고하시면 도움이 될거예요. end-to-end을 실현하려면 데이터가 많이 필요한 것이 핵심인데, 유넷은 이 부분을 augmentation으로 해결했다는 거죠. 어떤 증강기술을 사용했는지는 차차 설명하게 됩니다.

https://youtu.be/ImUoubi_t7s?si=JYQRmitCcbqp2CsD

Moreover, the network is fast. Segmentation of a 512x512 image takes less than a second on a recent GPU. 

게다가 처리속도도 아주 빠르다고 합니다.

여기까지 논문의 초록을 살펴봤는데요. 저는 초록에서 강조하고 있는 키워드가 아래와 같은 것으로 정리했습니다.

• Data Augmentation
• Contracting path & Expanding path
• end-to-end 
• Not a sliding window
• Outperform, fast

그럼 본격적으로 본문을 살펴보면서 위와 같은 키워드들이 어떻게 설명되는지 살펴보도록 하겠습니다.

Introduction

The typical use of convolutional networks is on classification tasks, where the output to an image is a single class label. However, in many visual tasks, especially in biomedical image processing, the desired output should include localization, i.e., a class label is supposed to be assigned to each pixel.

전형적인 CNN 모델은 이미지를 입력하면, 이미지에 해당하는 하나의 class label이 출력되는 방식으로 '분류'작업을 수행합니다. 그런데 실제로 많은 과업에서 이렇게 이미지 전체에 해당하는 하나의 레이블이 필요한 경우보다는, localization을 포함한 결과가 필요한 경우가 많습니다. 즉, 이미지 전체가 아닌 이미지 내부의 각 픽셀에 해당하는 labeling이 필요한 거죠.

위의 예시에서 맨 왼쪽은 "이 사진은 고양이 사진이다!"라고 전체에 대한 레이블링을 수행하였습니다. 하지만 localization은 모든 픽셀에 대해서 레이블링을 수행합니다. 그래서 그 옆의 그림처럼 이 사진의 어떤 부분은 고양이고, 또 다른 부분은 잔디이고, 이렇게 다양한 결과값을 출력할 수가 있는 것입니다.

Moreover, thousands of training images are usually beyond reach in biomedical tasks.

특히 바이오메디컬(생의학) 분야에서는 데이터 수가 부족한 경우가 많다는 것도 고려할 필요가 있는데요. 레이블링을 의사같은 전문가가 해야하기 때문에, dog-cat 분류같이 아무나 할수 있는게 아니라는 이유에서 그렇다고 합니다. 대회에 참가해서 30장 데이터로 승부를 봐야 했던 유넷처럼 말이죠......... (아니 근데 아무리 그래도 30장은 좀 너무했다)

Hence, Ciresan et al. [1] trained a network in a sliding-window setup to predict the class label of each pixel by providing a local region (patch) around that pixel as input. First, this network can localize. Secondly, the training data in terms of patches is much larger than the number of training images. The resulting network won the EM segmentation challenge at ISBI 2012 by a large margin.

유넷 이전 다수의 모델은 sliding window를 적용했습니다. sliding window를 적용하면 하나의 이미지로부터 여러 개의 patch가 생기면서 데이터 수 부족 현상을 해결할 수 있었기 때문이죠.

Obviously, the strategy in Ciresan et al. [1] has two drawbacks. First, it is quite slow because the network must be run separately for each patch, and there is a lot of redundancy due to overlapping patches. Secondly, there is a trade-off between localization accuracy and the use of context. Larger patches require more max-pooling layers that reduce the localization accuracy, while small patches allow the network to see only little context. More recent approaches [11,4] proposed a classifier output that takes into account the features from multiple layers. Good localization and the use of context are possible at the same time.

그런데 sliding window를 적용하면 계산량이 많아지면서 전체 네트워크가 느려진다는 단점이 있습니다. 게다가, localization accuracy와 use of context - 즉 전체와 부분 사이의 tradeoff가 발생하는 문제점도 있었는데요.

드디어 그 유명한 유넷의 구조가 언급됩니다. 

In this paper, we build upon a more elegant architecture, the so-called “fully convolutional network” [9]. We modify and extend this architecture such that it works with very few training images and yields more precise segmentations; see Figure 1. 

유넷 페이퍼에서는 이러한 문제점들을 해결하기 위해 U모양의 "fully convolutional network"를 제안합니다. 적은 수의 트레이닝 이미지 데이터로도 학습이 잘 되고, 정확한 segmentation이 가능하도록 말이죠.

The main idea in [9] is to supplement a usual contracting network by successive layers, where pooling operators are replaced by upsampling operators. Hence, these layers increase the resolution of the output. 

핵심 아이디어는 초록에서 살펴본 바와 같이 contracting - expanding network구조를 사용하는 것입니다. 입력된 이미지는 contracting network의 연속적인 레이어 구조를 거치며 사이즈는 줄어들고 채널 수는 늘어납니다. 이어서 expanding network의 연속적인 레이어 구조에서 pooling 대신 upsampling을 거치며 사이즈가 다시 커지고, 채널 수는 줄어드는데요. 

In order to localize, high resolution features from the contracting path are combined with the upsampled output. A successive convolution layer can then learn to assemble a more precise output based on this information.

One important modification in our architecture is that in the upsampling part we have also a large number of feature channels, which allow the network to propagate context information to higher resolution layers. As a consequence, the expansive path is more or less symmetric to the contracting path, and yields a 
u-shaped architecture.

Feature 1 이미지에서 회색 화살표를 보시면 대칭 구조를 이루며 레이어들이 서로 짝을 이루고 있습니다. contracting path에서 추출한 Feature map의 정보가 업샘플링 과정에서 skip-connection 방식으로 결합되는 과정을 도식화한 것입니다.

논문에서 'skip-connection'이라는 용어가 직접적으로 언급되지는 않지만, 이러한 방식을 많은 사람들이 그렇게 명명하고 있습니다. 스킵 연결방식을 사용해서 expanding path의 콘볼루션 레이어는 contracting path로부터 정보를 제공받게 되고, 보다 정확한 출력을 조합하는 방법을 학습할 수 있게 됩니다. 깊은 레이어 구조를 거치며 다양한 정보가 손실될 가능성이 있는데, 그런 점을 skip-connection을 사용하여 보완해 학습이 잘 되도록 구조를 설계한 것이라고 보시면 되겠습니다.

The network does not have any fully connected layers and only uses the valid part of each convolution
, i.e., the segmentation map only contains the pixels, for which the full context is available in the input image. This strategy allows the seamless segmentation of arbitrarily large images by an overlap-tile strategy (see Figure 2). 

유넷 아키텍처는 기본적으로 fully connect layers로 이루어져 있지 않습니다. 대신, 앞에서 살펴본 바와 같이 모든 레이어가 Convolution(fully convolutional layers)으로 이루어져 있어요. 또 유넷 네트워크는 각 컨볼루션 연산에서 유효한 부분(valid part)만을 사용합니다. 이는 입력 이미지의 가장자리를 패딩 없이 처리하는 방식을 의미하는데요, 따라서 결과적으로 출력 이미지의 크기가 입력 이미지보다 작아지게 되고, 모든 출력 픽셀은 입력 이미지의 전체 컨텍스트(문맥 정보)를 갖게 됩니다.

저는 여기서 처음에 들었던 의문이.... '패딩 없이 처리'를 한다고 해서 어떻게 전체 컨텍스트 정보를 갖게 된다는 것인가? 라는 거였어요. 패딩을 사용 안하면 안하는거지, 그걸 굳이 저렇게 거창하게 말할 필요가 있나.. 싶었거든요. 이 부분은 바로 이어서 언급되는 '오버랩 타일 전략'을 살펴보면서 간단히 해결할 수 있었습니다.

(+) 그런데 실제로 코드로 구현하는 경우 계산의 편의를 위해 valid padding 대신 same padding을 사용하는 경우가 많은 것 같습니다.

유넷은 이미지의 가장자리에 패딩을 사용하지 않는 대신, 오버랩-타일(overlap-tile) 전략이라는 걸 사용합니다. 이 방식은 특히 큰 이미지를 분할할 때 유용한데요. 경계 영역에서의 정보 손실을 최소화할 수 있는 장점이 있다고 언급하고 있습니다. 아래 그림을 함께 보시죠.

To predict the pixels in the border region of the image, the missing context is extrapolated by mirroring the input image. This tiling strategy is important to apply the network to large images, since otherwise the resolution would be limited by the GPU memory.

U-Net에서 설명하는 타일링 전략(overlap-tile strategy)은 하나의 큰 이미지를 여러 개의 작은 타일로 나누어 처리하는 방식입니다. 이 때 경계 부분에서 패딩을 사용하는 대신, 입력 이미지를 반사하는 미러링(mirroring)하는 방법을 사용하고, 타일끼리 서로 겹치는 부분이 발생하기 때문에 '오버랩'이라는 단어가 앞에 붙게 되었어요. 이 오버랩 타일링 전략을 통해 유넷은 이미지 경계에서 패딩을 사용했을때 발생할 수 있는 정보 손실을 최소화할 수 있었고, 덕분에 정확한 경계 예측을 할 수가 있었어요.

생각해 보면 segmentation을 수행할 때 경계값을 명확하게 파악할 필요성이 있다고 납득이 되더라고요. 이 타일링 전략은, 앞에 서론에서 잠깐 살펴보았던 슬라이딩 윈도우(sliding window)의 문제점을 해결한 전략이라고 할 수 있겠습니다. 이러한 타일링 전략은 특히 큰 고해상도의 이미지에 유넷 네트워크를 적용할 때 중요했다고 하는데, 그렇지 않으면 해상도가 GPU 메모리에 의해 제한될 수가 있었다고 하네요.

https://youtu.be/O_7mR4H9WLk?si=9aNrE8Olud-7bq-Z

슬라이딩 윈도우와 타일링 전략, 미러링 부분에 추가 보충이 필요하신 분들께서는 위의 유투브 영상을 참고하시는 것을 추천합니다. 시각적으로 설명이 깔끔하게 잘 되어 있어서 저도 도움을 많이 받았습니다 :)

As for our tasks there is very little training data available, we use excessive data augmentation by applying elastic deformations to the available training images. This allows the network to learn invariance to such deformations, without the need to see these transformations in the annotated image corpus.

U-Net은 고작 30장밖에 되지 않는 데이터셋으로 성능을 뽑아야 했기 때문에, augmentation을 빡세게 굴렸다! 라고 밝히고 있습니다. 특히 그 중에서도 elastic deformation을 사용했다고 밝히고 있어요.

Pixel이 랜덤하게 다른 방향으로 뒤틀리도록 변형하는 'elastic deformation'을 통해 자연스럽고 현실세계에 있을법한 새로운 데이터를 만들어내는 의의가 있었다고 해요. 위의 그림을 보면 elastic deformation을 거쳤을 때 기본 이미지의 특성은 어느정도 유지되면서 이미지가  다양하게 변형되는 것을 볼 수 있습니다. 논문에 따르면 이런 방법의 변형이 특히 biomedical 생의학 분야에서 의미있다는 연구 결과가 있었다며 참고 문헌을 밝히고 있네요. (세포라는 게 아무래도 원래 찌글짜글 거리니까…..)

Another challenge in many cell segmentation tasks is the separation of touching objects of the same class; see Figure 3. To this end, we propose the use of a weighted loss, where the separating background labels between touching cells obtain a large weight in the loss function. The resulting network is applicable to various biomedical segmentation problems. In this paper, we show results on the segmentation of neuronal structures in EM stacks (an ongoing competition started at ISBI 2012), where we outperformed the network of Ciresan et al. [1]. Furthermore, we show results for cell segmentation in light microscopy images from the ISBI cell tracking challenge 2015. Here we won with a large margin on the two most challenging 2D transmitted light datasets.

세포 분할 작업에서 해결해야 하는 또 다른 도전 과제가 있는데요. 바로 같은 클래스의 서로 접촉하는 객체를 분리하는 것입니다. 위의 그림을 보면, 여러 개의 같은 세포가 서로 매우 가깝게 접촉해 있는 것을 확인할 수 있습니다. 얘네가 서로 다른 개체임에도 불구하고 label이 같아서 하나의 큰 뭉텅이로 간주될 위험성이 있단 말이죠.

유넷은 이러한 문제점을 해결하고자 접촉하는 세포 사이의 배경 레이블에 큰 가중치를 부여하는 weighted loss를 제안합니다. 탐지된 객체가 서로 가까이 붙어 있을 수록 그 사이의 배경이 큰 weight값을 갖게 되는 원리인데요. 위의 Fig. 3. 그림에서 d를 보시면 되고, 세포 사이의 거리가 가까울수록 배경이 빨간 값을 가지게 되는 점에 주목하시면 되겠습니다. 이 방법을 적용하면서 뛰어난 segmentation 성능을 확보하게 되었고, 이 덕분에 대회를 우승할 수 있었다고 밝히고 있었습니다.

Network Architecture

다음으로 유넷 네트워크의 구조에 대해서 조금 더 자세히 살펴보도록 하겠습니다. 이 부분은, 유넷 그림의 구조를 보면서 스스로 설명할 수 있어야 합니다.

Contracting path (left side)

The contracting path follows the typical architecture of a convolutional network. It consists of the repeated application of two 3x3 convolutions (unpadded convolutions), each followed by a rectified linear unit (ReLU) and a 2x2 max pooling operation with stride 2 for downsampling. At each downsampling step we double the number of feature channels.

Contracting path는 기본적인 콘볼루션 네트워크 구조로 되어 있습니다. (3x3) 필터를 사용해서 패딩 없이 콘볼루션을 계산한 뒤 ReLU를 걸어주는 과정을 2번 반복하고, (2x2) 사이즈와 stride 2의 max poolng을 적용하여 다운샘플링을 진행합니다. 이렇게 한 셋트의 다운샘플링을 진행할 때마다 채널의 수는 2배씩 증가하게 됩니다. 그림에선 다운샘플링을 4번 했네요.

이렇게 이미지의 사이즈가 줄어드는 과정에서 이미지의 특징이 추출되고, 채널이 늘어나는 과정에서 이미지의 다양한 특징을 확보하게 됩니다.

Expansive path (right side)

Every step in the expansive path consists of an upsampling of the feature map followed by a 2x2 convolution (“up-convolution”) that halves the number of feature channels, a concatenation with the correspondingly cropped feature map from the contracting path, and two 3x3 convolutions, each followed by a ReLU. The cropping is necessary due to the loss of border pixels in every convolution.

Expansive path에서는 (2x2) 필터 사이즈의 up-convolution을 통해 특징 맵(feature map)의 사이즈를 키우는 업샘플링(upsampling)을 수행하게 되는데, 이 과정에서 반대로 채널의 수는 다시 절반으로 점차 줄어들게 됩니다.

이후에 contracting path에서 해당 단계의 크롭된(cropped) 특징 맵과 연결(concatenation)한다는 이야기가 나오는데요. 이 부분은 구조도에서 회색 화살표를 확인하시면 됩니다. 위의 그림에서 빨간색 표시한 부분을 잘 보면, Contracting path의 (64 * 64 * 512)이 copy and crop으로 Extracting path의 (56 * 56 * 512)와 concatenate되어 (56 * 56 *  1024)가 됩니다. 64와 56이 사이즈가 맞지 않으니 64를 56으로 crop해주어야 했을 것이고, 이후 필터 512장과 512장이 concat으로 이어붙어 1024장이 된 것이죠.
논문에서는 언급되지 않은 용어이지만 흔히들 이 과정을 Skip-connection이라고 부르는데, element-wise summation이 아닌 concatenate를 사용하는 이유에 대해서는 아래 포스팅을 참고해보시면 도움이 되실 거예요.

https://at0z.tistory.com/164

저는 여기서 앞서 살펴본 '미러링 타일 전략'과 개념 혼돈이 오면서 헷갈리는 시기가 있었습니다. upsampling과 downsampling 과정에서 콘볼루션을 계산할 때, 패딩을 사용하는 대신 미러링을 사용하는 것인가? 라는 오개념이 자리잡은 것인데요. 저와 비슷한 분을 위해 짚고 넘어가자면, 미러링은 이미지를 여러 개의 타일로 나누고 난 다음에, 그 각각의 타일 테두리에 미러링을 추가해 주어 경계 정보를 보완해주는 것으로, 그렇게 완성된 타일 이미지에 upsampling과 downsampling의 u-net 구조가 적용되는 것입니다. 즉, 

1. 전체 이미지를 작은 타일로 나누어서 사용한다
2. 이때 각각의 작은 타일의 가장자리에는 미러링이 적용된다
3. 이렇게 미러링된 하나의 패치 이미지에 각각 u넷 구조가 적용된다

이렇게 순서를 이해하시면 되겠습니다 :)

 At the final layer a 1x1 convolution is used to map each 64- component feature vector to the desired number of classes. In total the network has 23 convolutional layers.

Skip-connection 이후에 두 번의 3x3 컨볼루션과 각각의 ReLU 활성화 함수를 적용하게 되는데요. 전체 네트워크는 총 23개의 컨볼루션 레이어로 구성되고, Expanding path의 마지막 레이어에서는 1x1 컨볼루션을 사용하여 결과값의 필터 수를 원하는 값으로 조정합니다. 

To allow a seamless tiling of the output segmentation map (see Figure 2), it is important to select the input tile size such that all 2x2 max-pooling operations are applied to a layer with an even x- and y-size.

seamless tiling을 위해서는 입력 타일 크기를 신중하게 선택할 필요가 있는데, 특히 모든 2x2 max-pooling 연산이 x와 y 크기가 짝수인 레이어에 적용되도록 해야 한다고 밝히고 있습니다. 이는 각 풀링 연산이 끝나는 레이어의 크기가 짝수여야 다음 연산에서도 크기가 정확히 맞아 떨어지기 때문인데요.

논문에서 표현하는 'seamless tiling'이란, U-Net이 큰 이미지를 처리할 때 '타일링 전략'을 사용해서 하나의 이미지를 여러 개의 작은 타일로 나누었던 것을, 마지막에 다시 이어 붙이는 작업이 매끄럽고 이쁘게 잘 되는 것을 의미한다고 보시면 돼요. 마치 원래부터 하나의 이미지였던 것처럼 이쁘게 잘 이어지기 위해서는 각 풀링 연산이 끝나는 레이어의 크기가 짝수가 되도록 잘 맞춰주라! 라는 겁니다. 

Training

The input images and their corresponding segmentation maps are used to train the network with the stochastic gradient descent implementation of Caffe [6].

• 인풋 이미지와 그 이미지의 segmentation map 이미지가 네트워크 학습 데이터로 사용되었다고 합니다.
• SGD(stochastic gradient descent)를 이용하여 학습을 진행했습니다.

Due to the unpadded convolutions, the output image is smaller than the input by a constant border width. To minimize the overhead and make maximum use of the GPU memory, we favor large input tiles over a large batch size and hence reduce the batch to a single image. Accordingly we use a high momentum (0.99) such that a large number of the previously seen training samples determine the update in the current optimization step.

 U-Net은 '타일링 전략'을 사용한다고 앞서 밝혔는데요. 콘볼루션을 계산할 때 패딩 없이 진행이 되기 때문에 이미지가 너무 작은 경우 경계값 정보가 많이 손실되어 학습이 제대로 이루어지지 않을 위험성을 가지고 있습니다. 따라서 여러 개의 작은 타일로 나누는 것보다 타일 크기를 크게 해서 전체 타일 수가 적은 것을 선호한다고 밝히고 있습니다.

그런데 큰 타일 크기를 사용하면 그만큼 SGD(Stochastic Gradient Descent) 배치 수는 줄어야 합니다. GPU 메모리가 한정되어 있기 때문이죠. 배치 수가 줄어든다면 그만큼 학습할 수 있는 데이터 다양성도 함께 줄어드는 단점이 있습니다. 다르게 말하자면 경사 하강법 과정에서 그래디언트 변동성이 커질 수 있게 되고, 이는 최적화 과정을 불안정하게 만들 수가 있는 것이죠.

따라서 논문에서는 이러한 문제점을 해결하기 위해 높은 모멘텀 momentum(0.99)을 사용한다고 밝히고 있습니다. 높은 모멘텀을 사용하면 이러한 변동성을 완화하고, 모델이 더 안정적으로 최적화를 할 수 있기 때문인데요. 즉 큰 타일 크기 사용으로 인한 불안정성을 높은 모멘텀이 보완해주는 효과가 있는 것입니다.

모멘텀은 Local Minimum에 빠지는 경우를 대처할 수 있다는 특징이 있으며, 높은 모멘텀은 현재 그래디언트 업데이트에 이전 단계들의 그래디언트 영향을 더 많이 반영합니다. 모멘텀에 대한 더 깊은 이해가 필요하시면, 아래 포스팅을 참고하시면 도움이 될 것입니다.

https://heytech.tistory.com/382

Training에 사용한 U-Net의 Loss Function은 각 픽셀에 대한 에너지 함수(E)의 총합으로 구성되는데요. 에너지 함수는 다음과 같은 과정으로 계산됩니다.

1. 먼저 각 픽셀의 예측값에 소프트맥스(SoftMax) 함수 P를 적용해 예측값을 확률값으로 변환해 줍니다. 여기서 𝑙은 특정 클래스의 레이블을 의미합니다.
2. 그런 다음, 특정 클래스 𝑙에 대한 로그 확률을 계산합니다. 예측 확률은 0과 1 사이의 값을 가지게 되므로, 예측 확률이 1에 가까워질수록 로그 값은 0에 가까워집니다. 반대로 예측 확률이 0에 가까울수록 로그 값은 매우 큰 음수가 되겠죠.
3. 이제 각 픽셀의 실제 클래스 레이블과 예측된 확률 값을 비교하여 손실 값을 계산합니다. 이를 위해 교차 엔트로피 손실(Cross-Entropy Loss)을 사용합니다. 교차 엔트로피 손실 함수, 즉 '에너지 함수(E)'는 위와 같은데요. 특이하게 크로스 엔트로피에 픽셀 고유의 weight 𝑤(𝑥)을 곱함으로써 픽셀의 Loss값을 계산하는 것을 볼 수가 있습니다. 
4. 𝑤(𝑥) 식은 아래 수식과 같이 나타냅니다. 경계선(border)라인에 더 강한 학습을 시키기 위해서 가우시안 분포(Gaussian distribution)을 가정하고 경계선 근처의 픽셀에 더 높은 가중치를 부여하여 학습이 집중되도록 했다고 합니다.

마지막으로 모든 픽의 에너지 함수(E)를 합산하여 전체 이미지에 대한 손실을 계산합니다.

마지막 부분에는 초기 가중치 설정의 중요성을 강조하며, 가중치를 Gaussian 분포로부터 표준편차가 루트 2/𝑁인 값으로 초기 설정하였다고 합니다. 예를 들어, 3x3 convolution과 이전 레이어에 64개의 feature 채널이 있는 경우, 𝑁은  9⋅64 = 576이 되는 거죠. 이 부분은 가볍게 읽고 지나가겠습니다.

3.1. Data Augmentation

Data augmentation is essential to teach the network the desired invariance and robustness properties, when only few training samples are available. In case of microscopical images we primarily need shift and rotation invariance as well as robustness to deformations and gray value variations. Especially random elastic deformations of the training samples seem to be the key concept to train a segmentation network with very few annotated images. We generate smooth deformations using random displacement vectors on a coarse 3 by 3 grid. The displacements are sampled from a Gaussian distribution with 10 pixels standard deviation. Per-pixel displacements are then computed using bicubic interpolation. Drop-out layers at the end of the contracting path perform further implicit data augmentation.

논문에서 언급한 data augmentation 방식은 총 4가지 입니다.

• Shift
• Rotation
• Gray value
• Elastic Deformation

이 중에서도 Elastic Deformation이 key 역할을 했다고 밝히고 있습니다. 이 부분은 앞에서 짚고 넘어간 바가 있으니 추가 설명은 생략하겠습니다.

Experiments

UNet의 우수한 성능으로 각종 대회에서 최고 결과를 도출했다는 파트로, 이부분은 자세한 정리를 생략하겠습니다.

Conclusion

드디어! 마지막 결론입니다.

The u-net architecture achieves very good performance on very different biomedical segmentation applications. Thanks to data augmentation with elastic deformations, it only needs very few annotated images and has a very reasonable training time of only 10 hours on a NVidia Titan GPU (6 GB). We provide the full Caffe[6]-based implementation and the trained networks. We are sure that the u-net architecture can be applied easily to many more tasks.

U-Net 아키텍처는 매우 다양한 생의학적 세그멘테이션 응용에서 매우 우수한 성능을 발휘하는데, Elastic deformation을 이용한 Data Augmentation형을 이용한 데이터 증강 덕분에 가능했다고 밝히고 있습니다. 유넷의 성능 확보에 정말 핵심적인 역할을 한 게 분명합니다.

마지막으로 U-Net 아키텍처는 더 많은 작업에 쉽게 적용될 수 있을 것이라고 자신감을 나타내며 논문이 마무리되는데, 실제로 U-Net은 Biomedical 분야뿐만 아니라 이미지 segmentation이 필요한 다양한 최신 모델에 두루 널리 쓰이며 사랑받고 있습니다.

본 논문 리뷰를 통해 다양한 최신 모델에서 자주 보이는 U-Net의 구조와 원리에 대해서 이해할 수 있어서 영광이었습니다 :-) 다음에도 좋은 논문 리뷰로 찾아뵙겠습니다. 감사합니다!