Detectron2 꼼꼼 정리 2탄

앞서 1탄에서 FPN과 RPN 의 기능에 대해서 잘 알아보았고 이제 마지막으로 ROI Head를 설명해보고 아키텍쳐에 대한 기본설명을 마치려고 한다. ROI Head에서는 FPN Featuremap(P2~P5) 그리고 Region proposal box 그리고  Ground truth box를 입력으로 사용하게된다. 앞서 RPN 에서는 피쳐맵을 p~p6까지 사용하였지만 ROI Head에서는 p6는 사용하지 않는다. 

ROI Head는 StandardROIHeads class로 정의 되어있고 아래와 같은 개략도를 가진다.

먼저 Proposal Box를 재샘플링하게 되고 샘플링한 박스를 통해 ROI pooling을 거치게 된다.

샘플링한 roi들을 각각 p2~p5 중 어떤 feature map에서 roi를 잘라야 하는지 선택해야한다. 이를 적용하는 규칙은 다음과 같다. 

assigned feature level: floor(4 + log2(sqrt(box_area) / 224))

 

224는 표준 box 크기 이고 예를들어 224x224인 roi의 경우엔 p4 feature map에 할당되게 된다. 

레벨 할당 함수는 detectron2/modeling/poolers.py의 assign_boxes_to_levels 에서 실행되게 된다.

 

이제 roi를 자르기 위한 방법인 Roi Pooling이다. Roi Pooling이란 Faster R-CNN에서 ROI 영역에 해당하는 부분만 max-pooling을 거쳐 feature map으로부터 고정된 길이의 저차원 벡터로 축소하는 단계를 말한다. 

이 과정을 거침으로써 모든 roi들이 H*W 크기로 동일하여 지고 동일한 크기의 ROI들을 FC layer로 전달하기 위함이다. 

Mask rcnn에서는 이 Roi pooling이 Roi Align 이라는 네트워크로 교체된다. ROI Pooling은 detection을 위한 모델이여서 아주 미세하게 까지 정확한 위치 정보가 중요하지 않았다. 따라서 ROI 좌표가 소수점 좌표에 존재하면 좌표를 반올림하여 이동시킨후 Pooling을 진행하였는데 Mask rcnn에서는 정확한 좌표에서의 픽셀 정보가 중요하기 때문에 이를 bilinear interpolation을 이용해서 위치 보간을 한 방법론인 ROI-align을 사용한다. 정확하게는 Detectron2에선 ROI-align이 약간 수정된 ROI-alignv2를 사용한다. 이제 이렇게 p2~p5에서 자른 Roi들은 [B, C, H, W] = [N × 배치 크기, 256, 7, 7] 의 크기를 갖는다. 기본적으로 한 배치 N에 대한 ROI수는 512이고 ROI크기는 7X7로 고정된 사이즈이다. 배치 사이즈를 조정하고 싶다면 cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = 128 이런식으로 정의 해주며 바꾸면 된다. 이러한 7X7 ROI들은 BOX-head와 Mask_head 네트워크에 전달된다. Mask_head에서는 14x14로 Upsamling되어 전달된다. 

따라서 box_head에선 7x7 mask_head에선 14x14의 Roi들을 시용한다.

 

 (box_head): FastRCNNConvFCHead(
      (flatten): Flatten(start_dim=1, end_dim=-1)
      (fc1): Linear(in_features=12544, out_features=1024, bias=True)
      (fc_relu1): ReLU()
      (fc2): Linear(in_features=1024, out_features=1024, bias=True)
      (fc_relu2): ReLU()
    )

 

 (mask_head): MaskRCNNConvUpsampleHead(
      (mask_fcn1): Conv2d(
        256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)
        (activation): ReLU()
      )
      (mask_fcn2): Conv2d(
        256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)
        (activation): ReLU()
      )
      (mask_fcn3): Conv2d(
        256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)
        (activation): ReLU()
      )
      (mask_fcn4): Conv2d(
        256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)
        (activation): ReLU()
      )
      (deconv): ConvTranspose2d(256, 256, kernel_size=(2, 2), stride=(2, 2))
      (deconv_relu): ReLU()
      (predictor): Conv2d(256, 4, kernel_size=(1, 1), stride=(1, 1))

위의 box_head를 보면 [B, 256, 7, 7] 였던 입력텐서가 12544채널로 Flatten하여 fc layer에 공급되는걸 볼 수 있다.

최종적으로 아래와 같이 box_predictor 레이어인 cls_score 및 bbox_pred를 도출한다.

 

 (box_predictor): FastRCNNOutputLayers(
      (cls_score): Linear(in_features=1024, out_features=5, bias=True)
      (bbox_pred): Linear(in_features=1024, out_features=16, bias=True)
    )

 

  cls_score -> scores # shape:(B, class number+1)
  bbox_pred -> prediction_deltas # shape:(B, class number×4)

 

이제 훈련 loss를 계산하게 된다. 기본적으로 Mask rcnn의 loss는

  total_loss =  loss_box_reg + loss_cls + loss_mask + loss_rpn_cls + loss_rpn_loc 총 5가지 loss의 합으로 계산된다.

 

먼저 RPN network에서 loss_rpn_cls 와 loss_rpn_loc 를 계산 했었다. 간략하게 짚어보면

loss_rpn_cls 는 l1 loss이고 배경 그리드 포인트를 제외한 foreground인 grid point에서만 계산되는 손실이였다. 그리고

loss_rpn_loc 는 ground-truth objectness=1(foreground) 또는 0(backgraound)인 grid 지점에서만 계산되며 Binary cross entropy loss로 계산되었다. 

l1 loss 

loss_box_reg 는 bounding box의 위치에 대한 loss로 l1 loss로 계산된다. 

또한 loss_cls 는 class에 대한 분류 손실로 Softmax cross entropy loss로 계산된다. 이때 class수 K는 Class number +1(배경) 이 된다. 

loss mask는 Cross-entropy loss로 class branch에서 출력한 K class에 대해서만 loss를 계산한다. 

5가지 loss의 합으로 총 loss를 계산하고 loss를 줄이는 방향으로 학습하게 된다.

길게 mask rcnn을 예시로 Detectron2의 mask rcnn 모델 구조에 대해서 알아보았고 다음 포스팅에서 자세히 코드를 어떻게 다루는지 얘기해보고자 한다.