Efficient Interactive Annotation of Segmentation Datasets with Polygon-RNN++
1. Introduction
이 논문에서는 기존의 Polygon-RNN 모델을 아래와 같이 개선하였습니다.
- 신경망 구조를 변경
- 강화학습을 이용한 더 좋은 학습 알고리즘
- 모델의 한계로 인한 출력 폴리곤 해상도를 Graph Neural Network로 높이는 방법
2. Related Work
3. Polygon-RNN++
기존의 Polygon-RNN과 마찬가지로 사람이 이미 원하는 물체의 주변을 바운딩 박스로 표시하여 주어졌다고 생각하고, 박스를 15%정도 키운 다음 크롭하여 사용하였습니다. CNN+RNN 구조를 사용하였으며, CNN은 이미지의 특징 추출을, RNN은 그것으로부터 폴리곤 점을 디코딩하는 역할을 담당합니다.
우선, 기존의 모델에서 첫번째 점을 유추하는 엔코더 부분을 다시 디자인하였습니다. 다음 첫번째 점과 이미지 특징이 RNN 디코더로 입력되면 RNN은 각 타임 스텝마다 visual attention을 계산하여 점을 계산하게 됩니다. 학습된 evaluator network는 디코더로부터 생성된 점 후보 중 가장 좋은 것을 선택하게끔 합니다. 마지막으로 그래프 신경망이 폴리곤을 재조정하고 추가 점들과 함께 더 높은 해상도를 생성해 냅니다.
3.1 Residual Encoder with Skip Connections
CNN에서 연속적으로 다운샘플링 레이어를 반복되면 출력 해상도에 영향이 있게 됩니다. 이를 해결하기 위해서 ResNet-50 구조를 수정하여 네트워크의 stride를 줄이고 dilation을 사용하도록 하였습니다. 이렇게 하면 receptive field가 줄어들지 않고서도 특징 맵의 해상도를 더 크게 만들 수 있습니다. 또한 average pooling과 FC 레이어를 제거하였습니다.
다음 skip-layer 구조를 추가하여 저수준의 정보와 높은 수준의 정보를 모두 참고할 수 있도록 하였습니다. 원래 모델에서 VGG 모델의 제일 위에서 다른 레이어듸 특징들을 붙여 다운샘플링된 skip-layer 구조를 사용하였으나, 여기에서는 그 대신 사용할 수 있는 가장 높은 해상도에서 skip-layer의 모든 출력을 붙여 사용하도록 하였습니다. 이렇게 나온 특징 맵에 컨볼루셔널 레이어와 max-pooling 연산을 조합하여 마지막 특징 맵을 생성하였습니다. 컨볼루셔널 필터 크기는 3x3을 사용하였고, batch normalization과 ReLU를 사용하였습니다. Skip-connection에서 서로 다른 차원을 가지는 곳에서는 bilinear 업샘플링을 사용하였습니다.
3.2 Recurrent Decorder
이전 연구와 마찬가지로 컨볼루셔널 LSTM이을 사용하였습니다. 여기에 attention 메커니즘을 추가하여 같은 네트워크에 첫번째 점을 유추할 수 있도록 하였습니다.
3x3크기와 각각 64, 16 채널을 사용하는 2개 레이어의 ConvLSTM 을 사용하고, batch norm을 매 타임 스텝마다 적용하였습니다. 하지만 스텝마다 평균/분산은 서로 공유하지 않았습니다. RNN은 $D \times D + 1$의 크기의 one-hot 인코딩의 출력을 내놓도록 하였고, $D=28$을 사용하였습니다. 마지막 차원은 폴리곤이 끝 부분을 나타내는 토큰입니다.
Attention Weighted Features: Attentin 메커니즘을 위해서 매 타임 스텝 $t$마다 가중치가 적용된 특징 맵을 계산합니다.
$$ \alpha_t = \text{softmax}(f_{att}(x, f_1(h_{1,t-1}), f_2(h_{2, t-1}))) \\ F_t = x \circ \alpha_t$$
여기서 $\circ$는 Hadamard product이고, $x$는 skip feature 텐서를 나타냅니다. $h_{1, t}, h_{2, t}$는 ConvLSTM의 두개 레이어에서 나온 은닉 상태 텐서(hidden state tensor)를 나타냅니다. $f_1, f_2$는 은닉 상태 텐서를 완전 연결 레이어(fully-connected layer)를 통해 $\mathbb{R}^{D \times D \times 128} $로 매핑하는 함수입니다. $f_{att}$는 입력을 모두 더한 것을 입력으로 하여 $D \times D$로 매핑하는 완전 연결 레이어입니다. 이것은 결국 위치마다 attention 가중치를 지정하는 역할을 합니다.
이렇게 만들어진 $F_t$에 2개 이전의 점 one-hot 인코딩된 $y_{t-1}, y_{t-2}$과 첫번째 점 $y_0$을 쌓은 뒤, RNN으로 입력됩니다.
First Vertex: 이전 연구에서는 첫번째 점은 특수 상황으로 별도의 네트워크를 통해 추정하게 했지만, 여기에서는 skip-layer 구조를 이용하여 한 네트워크에 가지를 쳐 2개의 $D \times D$ 레이어를 쌓도록 하였습니다. 이전 연구와 마찬가지로 한 레이어는 에지를, 두번째 레이어는 폴리곤의 점을 추정하는 역할을 합니다.
3.3 Training using Reinforcement Learning
이전 연구에서는 타임 스텝마다 cross entropy loss를 이용하였습니다. 하지만 여기에는 두가지 문제가 있는데, 하나는 점을 GT와 정확하게 유추하지 않았지만, 에지 위의 위치를 유추한 경우 MLE가 이를 잘못된 것으로 여기는 문제입니다. 두번째는 최적화의 목표가 실제 평가에 이용되는 것(IoU)과 다르다는 점입니다. 게다가 다음 타임 스텝을 학습할 때 유추된 것을 넣는 것이 아닌 GT 점을 넣어 학습을 시키기도 합니다.
이를 해결하기 위해, MLE는 초기화 단계에서만 사용하고, 폴리곤 유추 과정을 강화학습 문제로 먼저 푼 다음 파인튜닝하는 방법을 사용했습니다.
3.3.1 Problem formulation
리커런트 디코더를 순차적인 결정 에이전트로 생각해봤습니다. 인코더-디코더 구조의 파라미터 $\theta$는 다음 점(vertex) $v_t$를 선택하기 위한 정책(policy) $p_\theta$를 정의합니다. 시퀀스의 마지막에 보상(reward) $r$을 계산합니다. 보상은 폴리곤으로 만들어진 영역과 GT의 마스크 영역 $m$의 IoU로 계산됩니다. 이 보상의 기대값을 최대화 하는 로스 함수를 꾸몄습니다.
$$ L(\theta) = -\mathbb{E}_{v^s \sim p_\theta} [r(v^s, m)]$$
여기서 $v^s = (v^s_1, \cdots, v^s_T)$로 시각 $t$의 샘플링된 점입니다. $r=\text{IoU}(\text{mask}(v^s), m)$ 입니다.
3.3.2 Self-Critical Training with Policy Gradients
기대값의 그래디언트를 계산하기 위해 REINFORCE 트릭을 사용했습니다.
$$ \nabla L(\theta) = -\mathbb{E}_{v^s \sim p_\theta} [r(v^s, m), \nabla \log p_\theta (v^s)]$$
실제로는 Monte-Carlo 샘플링을 한번 이용하여 그래디언트를 계산합니다. 이 방법은 큰 분산과 높은 불안정성을 가지는데, 이를 해결하기 위해 self-critical 방법을 이용하여 테스트 타임에 inference 보상을 베이스라인으로 이용했습니다. 이에 따라 로스 함수는 다음과 같습니다.
$$ \nabla L(\theta) = -\mathbb{E}_{v^s \sim p_\theta} [(r(v^s, m) - r(\hat{v}^s, m)) \nabla \log p_\theta (v^s)]$$
$r(\hat{v}^s, m)$는 greedy decoding을 이용한 모델로터 얻은 보상입니다. 모델이 점을 탐색하는 랜덤성을 조절하기 위해서 파라미터 $\tau$를 두어 정책의 softmax에 적용하였습니다. 실험에서는 0.6을 사용하였습니다.
3.4 Evaluator Network
첫번째 점을 선택하는 것을 매우 중요한데, 이것이 RNN의 시작점이 되기 때문입니다. 특히 가려진 객체의 경우에는 더 중요한데, 이 때에는 시작 점이 가려진 부분으로부터 먼 곳에서 시작하는 것이 좋기 때문입니다. RNN에서는 중요하지 않은 시퀀스를 자르기 위해서 bea search를 사용하는 경우가 보통인데, 보통 가려진 부분이 더 센 에지를 갖는 경우가 많아 이 부분의 log probability가 높기 마련이기에 이 부분이 잘려나가지 않을 여지가 많습니다.
이를 해결하기 위해 evaluator network를 사용하여 여러 폴리곤 후보들 중 하나를 효과적으로 선택하도록 하였습니다. evaluator network는 skip feature와 ConvLSTM의 마지막 상태 텐서, 유추한 폴리곤을 입력으로 합니다. 입력을 받으면 GT와의 IOU를 기반으로 폴리곤의 품질을 유추합니다. 이 네트워크는 2개의 3x3 콘볼루셔널 레이어와 FC 레이어로 구성됩니다. RL을 학습하면서도 이 네트워크를 동시에 학습할 수 있겠지만, RL의 파인튜닝이 모두 수렴한 후에 이 네트워크를 따로 학습하는 방법을 선택하였습니다.
학습 중에 $L(\phi) = [p(\phi, v^s) - \text{IoU}(m_{v^s}, m)]^2$를 최소화하였습니다. $p$는 유추된 IoU를, $m_{v^s}$는 샘플링된 점들의 마스크를, $m$은 GT의 마스크를 나타냅니다.
Inference: 테스트 시에는 K개 상위 첫번째 점들을 모아서 beam-search릉 이용해 폴리곤을 생성하였습니다. 생성된 K개의 서로 다른 폴리곤을 evaluator network에 넣어 가장 좋은 폴리곤을 선택하였습니다. 실험에는 $K=5$를 사용하였습니다. Beam-search 대신, 매 타임 스텝마다 evaluator network를 사용하면 inference 시간이 느려집니다.
3.5 Upscaling with a Graph Nueral Network
지금까지의 모델은 $D=28$로 해상도가 한정되어있었습니다. 이제 이것을 Gated Graph Neural Network (GGNN)을 이용해서 생성된 폴리곤을 더 높은 해상도로 올리고자 합니다. RNN 디코더를 학습할 때에는 GT가 이미 낮은 해상도로 줄여져 있습니다. 때문에 동일한 위치에 놓인 점들은 미리 제거가 되어있습니다. 따라서 더 높은 해상도에서는 추가로 점이 더 필요합니다. 이는 폴리곤의 위치 정보를 변경하게 될 것입니다.
이 모델은 RNN디코더로부터 생성된 점들의 시퀀스를 입력으로 합니다. 이 점들은 그래프의 노드로 생각하고, 높은 해상도의 디테일을 만들기 위해 연속된 두 노드 사이에 하나의 노드를 추가하도록 하였습니다. 추가된 노드의 위치는 두 노드 사이의 에지에 놓일 것입니다. 또한 마지막과 첫 점은 연결하도록 할 것입니다.
GGNN은 RNN은 어떤 그래프로 확장한 다음, 각 노드를 출력 하기 전에, 노드 사이의 정보를 효과적으로 유추하도록 합니다.
3.6 Annot. New Domains via Online Fine-Tuning
4. Experiments Results