[논문읽기] Structured Knowledge Distillation for Semantic Segmentation
이 글은 2019 CVPR 논문, Structured Knowledge Distillation for Semantic Segmentation 를 참고하여 작성하였습니다.
Knowledge distillation은 실제로 서비스에 적용하는 측면에서 굉장히 중요한 부분인 것 같아서 관심이 많은 연구주제 중 하나입니다. knowledge distillation에 대해서는 Ujjwal Upadhyay’s Blog에 정의와 필요성에 대해서 잘 설명되어 있으니 참고하면 좋을 것 같습니다. (+ Distilling the Knowledge in a Neural Network 논문)
이 논문에서는 Semantic Segmentation을 위한 Knowledge distillation 방법을 제안합니다. 다음과 같은 3가지 distillation 방법으로 정리할 수 있을 것 같습니다.
Approach
일단, 이 논문에서 제안한 distillation framework 입니다.
이 구조를 간단히 설명하자면, 처음에 Teacher net/Student net에 RGB image가 input으로 들어가고, Feature map이 계산됩니다. (여기서 Teacher network/Student network는 각각 cumbersome network/compact network로도 불립니다.) 이 feature map들은 upsample/classifier를 통해 최종 segmentation map이 나오게 됩니다. 이때 학습을 위해 사용되는 것이 Pixel-wise loss, Pair-wise loss, Holistic loss 입니다. 각각에 대해 더 자세히 알아보겠습니다.
pixel-wise distillation
pixel-wise distillation은 cumbersome network에서 생성된 pixel에 해당하는 class probability를 compact network에 transfer 합니다. 이때 각각의 probability를 transfer 해야하기 때문에 두 확률분포의 차이를 계산하는 데에 사용하는 함수인 쿨백-라이블러 발산(Kullback–Leibler divergence, KLD)을 사용하여 loss를 계산합니다.
\[\mathcal{l_{pi}}(S) = {1 \over {W'\times H'}}\sum_{i \in R}KL(q_i^s||q_i^t)\]\(q_i^s\) compact network S에서 생성된 i번째 픽셀의 class probability이고, \(q_i^t\)는 cumbersome network T의 class probability 입니다. 여기서 R은 (1, … , W’xH’)입니다. 둘 사이의 KL을 계산한 것이 loss가 되는 것입니다.
pair-wise distillation
semantic segmentation은 structured prediction problem이므로 structure information을 compact network에 transfer 해야합니다. 이때 사용하는 방법이 pair-wise distillation과 밑에 설명드릴 holistic distillation 입니다. 여기서는 픽셀 사이의 pair-wise similarities를 계산합니다.
\[\mathcal{l_{pa}}(S) = {1 \over ({W'\times H'})^2}\sum_{i \in R}\sum_{i \in R}(a_{ij}^s - a_{ij}^t)^2\]여기서 \(a_{ij}\)는 \(f_i^Tf_j/({\lVert f_i \rVert}_2{\lVert f_j \rVert}_2)\) 식을 통해 구하는데, f는 해당하는 픽셀의 feature를 의미합니다. feature map의 demension을 (W,H,N)이라고 할때, f는 (1,N)입니다. 따라서 \(a_{ij}\)값은 scalar가 되겠죠.
holistic distillation
holistic distillation은 compact network와 cumbersome network의 segmentation map들 사이의 high-order relation을 위해 사용됩니다. 이때, conditional generative adversarial learning의 컨셉과 Wasserstein distance가 사용됩니다. input RGB image I가 주어졌을 때, student network에 의해 생성된 segmentation map \(Q^s\)는 fake, teacher network에 의해 생성된 \(Q^t\)는 real로 해서 둘 사이의 분포 차이를 계산합니다.
\[\mathcal{l_{ho}}(S,D) = \mathbb{E}_{Q^s \sim p_s(Q^s)}[D(Q^s|I)] - \mathbb{E}_{Q^t \sim p_t(Q^t)}[D(Q^t|I)]\]이대 segmentation map과 conditional RGB image는 concatenate 되어서 discriminator input으로 들어갑니다. discriminator는 input image와 segmenation map이 얼마나 잘 매치가 되는지 판단합니다.
Experiments
각 네트워크 구조는 다음과 같습니다.
| teacher network | student network |
|---|---|
| PSPNet with ResNet101 | ResNet18, MobileNetV2Plus, MobileNetV2 … |
데이터셋은 Cityscape, CamVid, ADE20K를 사용하였고, Evaluation Metric으로는 segmentation에서 많이 쓰이는 IoU, model size는 파라미터 수로 측정을 하고, Complexity는 FLOPs로 측정했습니다. 각각에 대한 결과는 논문에 자세히 나와있습니다.
이전 논문에 비해서는 수학적인 부분이 많이 어렵지 않아서 이해하기 쉬웠던 것 같습니다. conditional generative adversarial learning, WGAN에 대해 더 자세히 공부해보고 싶어졌습니다. (WGAN 언제 공부하지...)
Leave a comment