본문 바로가기

Vision/Segmentation

[detectron2] unet 적용 시 num_class 이슈 및 해결방법

 

Facebook에서 개발한 detectron2의 U-net Implement 라이브러리(아래 링크)를 활용하는 과정에서 문제가 발생

꽤 많은 삽질(약 이틀)을 하여 기록하기 위해 남겨둠

https://github.com/zongyue-lu/pytorch-unet-family

 

1. 이슈 요약

이슈는 크게 세가지로 나누어 볼 수 있음

  1. Semantic Segmentation용 데이터 생성 및 로드 관련 이슈
  2. NUM_CLASS 이슈
  3. Cross Entropy Loss 이슈

 

2. 이슈 상세

1) 데이터 생성 과정

- 기존 COCO object detection용 데이터를 SS용 데이터로 변환하는 과정에서,  각 픽셀에 label 정보를 입력하고 png 형태로 저장

- label 정보는 1부터 시작하여 C(Num or Class)까지 annotation되어 있기 때문에 자연스럽게 background 값이 0으로 맵핑됨

 

2) NUM_CLASS 이슈

- detectron2 플랫폼 및 pytorch 프레임워크 상 NUM_CLASS를 지정해주어야 하는 부분이 두군데 있음

  • cfg.MODEL.ROI_HEADS.NUM_CLASSES
  • cfg.MODEL.SEM_SEG_HEAD.NUM_CLASSES

 

- 여러 도큐먼트 및 구글링을 찾아보면, Semantic or Panoptic Segmentation에서 두개의 값을 다르게 하라는 경우가 있었음. SEM_SEG_HEAD의 클래스의 갯수는 C+1로 설정하라 등

https://github.com/facebookresearch/detectron2/issues/1201

 

3) Pytorch CE Losses

- 공식 문서 상, CE losses에서 (predictions, targets) 두 값을 필수로 입력해야 하는데,  predictions(input)의 tensor shape은 (N, C, d1, d2, ..., dk), targets(ground truth)의 tensor shape은 (N, d1, d2, ..., dk) 로 정의되어 있음.

https://pytorch.org/docs/stable/generated/torch.nn.CrossEntropyLoss.html

 

- 여기에서 N은 batch size, C는 number of classes, d1~dk는 input tensor의 dimension임

- 나의 경우, 512 x 512 x 3 크기의 이미지를 인풋으로 넣었으므로 d1 = 512, d2(k=2) = 512이고, batch size 2, num classes 5로, 최종적인 predictions의 shape은 (2, 5, 512, 512) 임

- 추가로, targets(ground truth)의 input shape은  (2, 512, 512) 이고, 여기에서 중요한 것은 label 정보가 [0, C)로 0 ~ C-1 사이의 값을 가져야 함. (이 부분이 1)번에서 생성한 데이터와 충돌이 발생된 부분으로 판단됨)

 

3. 이슈 해결

1) GT Label 정보 변경

- 모든 tensor에서 -1하여 0 ~ C-1 범위로 값을 조정.

- background는 100으로 임의로 설정 함 (관련하여, 8-bit 컬러 하에서 최대값이 255 까지인데, 그 이상 넘어가는 경우 어떻게 될지가 궁금해짐)

- 인풋되는 Mask format을 bitmask로 설정

  • cfg.INPUT.MASK_FORMAT = 'bitmask'

 

2) Loss Function 수정 (ignore_value 활용)

- 위에서 임의로 설정한 100을 ignore 하기 위해서 기본 sem_seg 코드의 losses 함수를 확인 (ignore_value = 100)

def losses(self, predictions, targets):
    predictions = predictions.float()  # https://github.com/pytorch/pytorch/issues/48163
    # print(predictions.shape)
    loss = F.cross_entropy(
        predictions, targets, reduction="mean", ignore_index=self.ignore_value
    )
    losses = {"loss_sem_seg": loss * self.loss_weight}
    return losses

 

3) Num class 세팅

- Number of classes가 5이고, 두 값을 다르게 세팅하는 것이 아니라 같게 세팅 하였음. (C는 background는 제외한 foreground만 카운트)

  • cfg.MODEL.ROI_HEADS.NUM_CLASSES = 5
  • cfg.MODEL.SEM_SEG_HEAD.NUM_CLASSES = 5

 

4. 결론

위와 같이 세팅하였을 때, 문제 없이 학습이 진행됨을 확인하였음.

일단은 임시방편으로 수정한 면이 있는데, 좀 더 근본적인 방법으로 해결 방법을 모색해야 함.

추론 결과 이상 여부는 추가 확인 필요.

 

 

끗.