U^2-Net:Going Deeper with Nested U-Structure for Salient Object Detection

Salient Object Detection

Featured image

Main Contribution

Network architecture $U^2$-$Net$을 제안했습니다. 이 네트워크는 2단계의 중첩된 U구조를 가지는데, Salient object detection을 위해 설계되었으며, 3가지의 장점을 가집니다.

  1. 이러한 구조는 다른 size의 RSU block의 receptive field를 혼합했기 때문에 다양한 scale로부터 더 많은 contextual information을 capture할 수 있습니다.
  2. RSU block 내부의 풀링 연산으로 인해 메모리와 계산 비용의 증가없이도 고해상도를 유지하면서 네트워크를 깊게 만들 수 있습니다.
  3. Classification 기반의 pre-trained 모델이 아니기 때문에 처음부터 학습시켜도 좋은 성능을 달성할 수 있습니다.

Architecture

Residual U-block

RSU structure.png local 및 global contextual information은 salient object detection에서 아주 중요하게 작용합니다. VGG나 ResNet, DenseNet과 같은 modern CNN의 설계에서, 1x1, 3x3과 같은 작은 convolution filter는 특징 추출에 자주 등장하는데, 이들은 저장 공간도 적게 차지하면서 계산적으로도 효율적이기 때문입니다. 그림의 a~c는 작은 filter를 사용하는 예시입니다. 이러한 shallow layer의 feature는 local feature만을 가지고 있습니다. 이는 filter가 global feature를 찾기에는 너무 작기 때문입니다. shallow layer가 고해상도의 feature맵에서 global information을 잘 얻기 위해서는 receptive field를 확대해야 합니다. 그림의 d는 inception과 같은 block인데 dilated convolution을 이용해 receptive filed를 확대하여 local과 global feature를 모두 추출할 수 있게 했습니다. 하지만 input과 같은 깊지 않은 layer에서 이러한 연산을 수행하면 메모리와 계산 비용이 많이 필요합니다. 본 논문에서는 위의 문제들을 해결해 intra-stage multi-scale feature를 찾기 위해 U-Net에서 영감을 받아 RSU(ReSidual U-block)을 제안하였습니다. 그림의 e에 $RSU-L(C_{in},M,C_{out})$의 구조가 나와있습니다. $C_{in}, C_{out}$은 input, output 채널의 수이며, M은 layer 내부의 채널 수 입니다. 따라서 RSU는 3가지 component로 구성되어 있습니다.

  1. Input convolution layer 입니다. 이는 input을 $C_{out}$채널을 가진 intermediate map $F_1(x)$로 변환시킵니다. 이는 local feature를 추출하는 plain convolutional layer입니다.
  2. 두번째로 U-Net과 비슷한 L 높이의 encoder-decoder 구조입니다. 이는 $F_1(x)$를 입력으로 받으며, multi-scale contextual information $U(F_1(x))$를 추출하고 encode 하기 위해 학습합니다. L의 크기가 클수록 더 많은 pooling 연산, 큰 범위의 receptive field 그리고 풍부한 local, global feature를 갖습니다. 이러한 parameter로 구성하는 것이 input feature map으로부터 더욱 multi-scale feature를 잘 추출할 수 있게 해줍니다. multi-scale feature는 점차적으로 downsampling된 feature map들로부터 추출되며, 점차적으로 upsampling, concatenation, convolution을 통해 고해상도의 feature map로 인코딩 됩니다. 이러한 과정은 large scale로의 upsampling으로 인해 발생할 수 있는 detail의 손실을 완화해줍니다.
  3. residual connection은 local feature와 multi-scale feature를 합쳐줍니다:$F_1(x)+U(F_1(x))$.

    Res vs RSU.png 원래의 ResNet 구조와 비교해보겠습니다. 원래의 residual block은 input feature $x$와 convolution layer $F_2, F_1$ 으로 이루어지는데, RSU의 경우 single-stream convolution을 U-Net like structure로 바꿨고, original feature를 local feature로 대체했습니다. 이러한 설계는 네트워크로 하여금 각 residual block에 multiple scale로부터 직접 feature를 추출하게 합니다. 또한, RSU는 대부분의 연산을 downsampled feature map에 수행하기 때문에 계산 오버헤드가 작습니다. computation cost.png 위 그림을 보면 DSE나 INC와 같은 block들은 채널 수에 비례해 계산량이 증가하는 반면, RSU는 그에 비하면 거의 차이가 없습니다.

$U^2$-$Net$

U^2-net structure.png $U^2$-$Net$은 위에서 설명한 RSU block을 사용하여 만든 네트워크로, 위 그림을 보면 U 구조를 가진 것을 알 수 있습니다. 이로 인해 중첩된 U구조를 갖게 되었고, 그렇기 때문에 $U^2$-$Net$이라 명명되었습니다. $U^2$-$Net$은 11개의 블록으로 구성되어있고, RSU로 구성된 만큼 intra-stage multi-scale feature들을 잘 추출할 수 있으며, 이러한 특징들을 더 효과적으로 aggregation 합니다. 위 그림을 보면 encoder와 decoder가 대칭으로 이루어져 있는데, 각 encoder와 decoder는 같은 RSU block으로 이루어져 있습니다. EN_1부터 4까지 RSU-7, RSU-6, RSU-5, RSU-4로 이루어져 있으며, 여기서 각 숫자는 위에서 설명드린 RSU-L의 L입니다. 또한, En_5, 6의 경우 RSU-4F로 구성되어 있는데, 여기서 F는 기존 RSU block의 pooling과 upsampling 연산을 dilated convolution으로 대체한 것을 의미합니다. 따라서 RSU-4F는 input feature map과 intermediate feature map의 해상도가 동일합니다. 마지막으로, 모든 side의 출력을 concatenation해 convolution layer를 거쳐 최종 saliency map $S_{fuse}$를 출력하게 됩니다.

Supervision

학습 과정에서 HED와 비슷한 deep supervision을 사용했습니다. 이것을 효과는 HED와 DSS에서 증명되었습니다. loss는 다음과 같습니다. $L = \sum_{m=1}^Mw_{side}^{(m)}l_{side}^{(m)} + w_{fuse}l_{fuse}$ 여기서 $M$은 위 그림 기준으로 6입니다. side의 수가 6이기 때문입니다. 또한, $w$는 각 loss의 weight이며, $l$은 표준 binary cross entropy 입니다.

Experiment results

result1.png 위에서 설명한 것과 같이 많은 데이터셋에서 좋은 성능을 내고 있는 것을 보실 수 있습니다.

result2.png 정량적인 성능 뿐만이 아니라 정성적인 성능 또한 뛰어난 것을 보실 수 있습니다.

Reference

$U^2$-$Net$ paper : https://arxiv.org/pdf/2005.09007.pdf