[논문읽기] REINFORCED ACTIVE LEARNING FOR IMAGE SEGMENTATION

1. 개요 및 연구동기

픽셀 단위 레이블을 획득하는 것은 비용과 시간이 많이 든다. 그리고 현실적인 segmentation 데이터세트는 매우 불균형하다. 이미지 분할 작업에서 강화학습을 사용한 새로운 active learning 전략을 제안한다. 

 

강화학습을 사용하는 현재의 AL 기술은 label budget 이 충족될때까지 단계당 하나의 샘플에 레이블을 지정한다. 이는 각 단계마다 segmentation 네트워크의 업데이트 및 보상이 포함되기 때문에 매우 비효율적이다. 

 

우리는 분류 네트워크가 최소한의 레이블링 된 데이터로 좋은 성능을 달성할 수 있도록 레이블을 요청할 데이터를 통해 policy를 학습하는 것을 목표로 한다. 전체 이미지 대신 영역을 선택하면 그림 1과 같이 알고리즘이 가장 관련성이 높은 부분에 집중할 수 있다. 

 

 

 

 

3. Method

우리는 segmentation 네트워크 $f(\theta)$ 의 성능을 최대화하기 위해 레이블이 지정되지 않은 대규모 데이터세트에서 소수의 영역을 선택하는데 관심이 있다. 이 프로세스는 주어진 budget $B$ 가 달성될 때까지 반복적으로 수행된다. 각 반복 $t$에서 쿼리 네트워크 $\pi(\phi)$ 는 레이블이 지정되지 않은 큰 데이터세트 $U_t$ 에서 레이블을 지정할 $K$개의 영역들을 선택한다.  이러한 샘플들은 segmentation network $f$를 훈련시키는데 사용되는 데이터세트 $L_t$에 추가된다. 

 

성능 측정은 IoU 로한다. 

 

우리는 AL 문제를 Markov Decision Process 로 공식화했다. 

 

우리는 query network $\pi(\phi)$ 를 DQN 으로 모델링했다. 이러한 접근방식을 통해 모델은 이전 AL 경험을 기반으로 선택전략을 학습할 수 있다. 

 

각 에피소드 $e$ 는 총 $T$단계를 거친다. 분할 네트워크 f 를 초기 가중치 세트 $\theta_0$ 세트로 설정하고 주석 데이터 없이 시작한다. 각 반복 t 에서 다음 단계가 수행된다.

그림 2는 학습 알고리즘을 설명한다. 

 

State representation

데이터의 하위집합 $D_S$에서 segmentation network $f_t$의 예측을 사용하여 state $s_t$의 전역표현을 생성한다. (그림 2의 1단계) 픽셀단위의 예측으로 인한 많은 양의 메모리 사용을 피하기 위해 간결한 representation 이 필요하다. $D_S$ 의 샘플은 패치로 분할되면 모든 샘플에 대해 compact feature vector 가 계산된다.. 각 영역은 두 가지 기능 세트의 연결로 인코딩된다. 하나는 class 예측을 기반으로 하고 다른 하나는 Shannon Entropy 라고 표시되는 uncertainty 를 기반으로 한다. 

 

features 의 첫번째 집합은 각 분할 클래스에 대해 예측된 픽셀 수의 count 이다. 이 기능은 작은 패치에 덜 중요한 공간 정보는 무시하면서 주어진 패치에 대한 분할 예측을 인코딩한다. 또한 예측 클래스의 확률에 대한 엔트로피로 예측 변수의 uncertainty 를 측정한다. 각 영역에 대해 엔트로피를 계산하여 공간 엔트로피 맵을 얻는다. 그리고 최소, 최대, 평균 풀링을 사용하여 다운 샘플링 된 feature map을 얻는다. 두번째 feature set 는 이러한 엔트로피 feature 들을 병합하고 연결하여 얻는다. 

 

마지막으로 state $s_t$ 는 $D_S$에서 각 영역의 feature representation 의 앙상블로 표현된다. 

Action Representation

이 작업에서 action 은 레이블이 지정되지 않은 영역에 주석을 요청하는 것을 의미한다. 

각 단계 t 에서 각각 균일하게 샘플링된 영역을 포함하는 레이블이 지정되지 않은 영역을 샘플링하여 전체 집합을 근사화한다. 

각 하위동작 $a^{k,n}_t$ 는 네 가지 다른 feature 의 concatenation 이다.

1. 엔트로피,

2. class distribution features, 

3. 영역 $x_k$ 와 레이블이 지정된 집합 간의 유사성, 

4. region 과 레이블이 지정되지 않은 집합 간의 유사성.

 

직관적으로 query network 는 레이블이 지정되지 않은 집합에서 대표 샘플을 계속 가져오면서 클래스 균형이 더 잘잡힌 레이블 집합을 구축하는 방법을 학습할 수 있어야 한다. 세분화 데이터세트의 클래스 불균형을 완화하고 전반적인 성능을 향상시키는데 도움이 될 수 있다. 

 

3.2 Batch Mode DQN

레이블이 지정된 세트 $D_T$ 로 DQN 을 교육하고 분할 $D_R$을 사용하여 rewards 를 계산한다. query agent 는 다음 state로 전환하기 전에  $K$개의 영역을 선택한다. 각 영역은 독립적으로 선택된다고 가정한다. 계산을 쉼게 하고 반복되는 영역 선택을 파하기 위해 다음과 같이 정의된 $P^k_t$ 에서 영역 $x_k$를 선택하도록 제한한다.

 

네트워크는 temporal difference(TD) error 기반으로 loss를 최적화하며 훈련된다.. 

훈련을 안정화하기 위해 double DQN을 사용하는 target network를 사용했다. action 은 target network 로 선택되고 query network 로 평가된다. 각 sub action 에 대한 TD 목표는 다음과 같이 표현된다. 

 

$\gamma$ 는 discount factor 이다.

이 공식은 하위 동작이 state 에 따라 서로 독립적이라는 가정 하에서 유효하다. 우리는 하위동작 K 의 수를 늘리는 것이 분할 성능에 부정적이지 않다는 것을 관찰했다. 부록 E.3 을 참고.

 

 

4. Experiments

4.1 Experimental setup

사용한 데이터셋 

CamVid, Cityscapes

 

Implementation details

모델은 5번 학습시켰다.(초기화시킨 상태의 모델이 5개)

보편적인 data augmentation 을 사용하였다.

 

Evaluation

 

 

 

4.2 Results

 

 

 

 

5. Conclusion

data, region  기반에서 이미지 분할을 위한 active learning 방법을 제안했다. 

목표는 루프에서 픽셀단위 레이블을 얻는 비용을 최소화하는 것이다. DQN 공식의 새로운 수정을 제안한다. 

클래스별 평균 IOU 를 직접 최적화하고 과소 표현된 클래스의 레이블을 더 많이 요구한다. 

 

 

 

참고자료

https://arxiv.org/abs/2002.06583

Reinforced active learning for image segmentation

https://github.com/ArantxaCasanova/ralis

GitHub - ArantxaCasanova/ralis: Code for the paper "Reinforced Active Learning for Image Segmentation"

 

코드 실행방법

1. 코드저장소 git clone 으로 다운받기

2. 데이터 다운받기 (깃허브 페이지에 링크가 달려있다)

3. pretrained segmentation model 다운받기

4. 소스코드 내에 있는 script 에 있는 sh 파일을 실행하는 것을 기본으로 한다.

5. sh 데이터파일 경로, pretrained 모델의 경로를 수정해준다.

6. data 폴더에서 데이터를 다루는 py 파일이 있는데 여기서도 path 를 바꿔주어야 한다. sh 파일에 있는 경로와 py 파일에 있는 경로를 concat 하여 최종 경로로 사용된다.

7. 쉘 스크립트 파일은 scripts 폴더가 아니라 ralis 폴더에서 실행되어야 정상작동한다.

8. camvid dataset 안에 있는 txt 파일은 없어도 작동한다.