[Experiment] Heteroscedastic Aleatoric Uncertainty Loss 코드 추가

[seonahmin]

Background

멘토님께서 말씀해주신 loss를 적용하면 segmentation prediction의 noise에 도움을 줄 수 있을 것 같아 코드를 만들어보고자함

Content

관련 논문 : What Uncertainties Do We Need in Bayesian Deep Learning for Computer Vision?

Details

[seonahmin]

이 loss를 사용하려면 model의 마지막에 input channel=(decoder의 마지막 channel)(=16)이고 out_channel이 1인 layer, 즉 Conv2d(in_channels=decoder_channels[-1], out_channels=1) 를 추가해야 하는 것 같습니다. 즉, 현재는 build_model로 만들어진 model이 입력을 받으면 mask를 return하는 방식으로 돌아가고 있는데 mask 뿐만 아니라 추가된 레이어에서 나오는 1개의 값을 추가로 return해야 한다는 것입니다. 하지만 이 과정에서 smp에서 제공하는 모델 사이의 layer를 추가하고 2개의 return값을 받도록 수정하는 법을 잘 모르겠어서 혹시 할 수 있을 것 같으신 분은 도와주시면 감사하겠습니다. smp모델을 사용하지 않고 처음부터 부캠에서 제공해 준 코드로 한다면 할 수 있을 것 같기는 하네요.

[Dongwoo-Im]

동일 논문인데 Figure 추가된 버전이 있어서 공유합니다. https://arxiv.org/pdf/1703.04977.pdf

[Dongwoo-Im]

잘 모르지만 custom head를 만들고 이걸 aux_params 처럼 동작하게 하면 될 것도 같습니다.

ref 1을 보면 smp 모델은 모두 aux_params라는 인자를 갖고 있습니다. ref 2를 보면 aux_params의 head로는 SegmentationHaed, ClassificationHead가 있습니다. 즉, ref 2에 Heteroscedastic Aleatoric Head를 추가로 구현하고 이를 모델에 적용한다면 될 것도 같습니다.

다만 ref 3 94번째 줄을 보면 aux_params를 추가했다고 해서 output이 달라지는 게 아니기 때문에, custom head 에서 모든 작업이 끝나야 할 것 같습니다.

ref 1 : https://smp.readthedocs.io/en/latest/_modules/segmentation_models_pytorch/decoders/unet/model.html#Unet ref 2 : https://github.com/qubvel/segmentation_models.pytorch/blob/740dab561ccf54a9ae4bb5bda3b8b18df3790025/segmentation_models_pytorch/base/heads.py#L13 ref 3 : https://github.com/Hyerin-oh/Naver_AI_Boostcamp_Pstage/blob/main/Pstage3_SemanticSegmentation/src/train.py

[seonahmin]

감사합니다 동우님! 동우님이 주신 ref 2와, 추가적인 ref 4 에 따르면 aux_param값이 존재하면 encoder에서 출력된 output이 classification head의 입력으로 들어가서 적용되는 루트가 추가되는 것 같습니다. SegmentationHead는 decoder의 output을 입력으로 받구요. ref 3에서는 aux_param을 어떻게 사용하는 지를 알 수 있는 것 같은데 사용한 것을 보면 단순히 classificationHead의 parameter값만 dict형태로 전달해주는 형태인 것 같습니다. 제가 모르겠는 점은 제가 Heteroscedastic Aleatoric Head를 추가로 구현했다는 가정 하에 이를 어떻게 코드에 접근해서 추가하는지 입니다. 제가 하고자 하는 바를 정리하자면, decoder 바로 뒤에 segmentation head와 Heteroscedastic Aleatoric Head를 각각 붙여서 2개의 output을 출력하고 싶습니다. decoder output 뒤에 -> output 1 : segmentation head output -> output 2: Heteroscedastic Aleatoric Head output
즉, segmentation head와 hetero head는 동시에 decoder의 output을 입력으로 받는셈이죠 쉽게 설명하고 싶은데 어렵네요ㅜㅜ

[Dongwoo-Im]

1. ref 2에 aleatoric_head가 구현되어 있다는 가정 하에, 올려주신 ref 4에서 forward 함수를 아래처럼 수정하면 output도 control 할 수 있을 것 같습니다.
2. 질문하신 ref 3의 aux_params는 결국 지금 구현되어 있는 aux_params의 사용 예시일 뿐, aleatoric_head를 구현한다면 이에 필요한 aleratoric_params를 사용할 것이고, 어떤 parameter를 받든지 상관없다고 생각합니다.
3. 아래는 지금 생각해본 하나의 예시일뿐, new_masks = self.aleatoric_head(masks) 처럼 구현해도 될 것 같긴 합니다.

def forward(self, x):
        """Sequentially pass `x` trough model`s encoder, decoder and heads"""

        self.check_input_shape(x)

        features = self.encoder(x)
        decoder_output = self.decoder(*features)

        masks = self.segmentation_head(decoder_output)

        if self.classification_head is not None:
            labels = self.classification_head(features[-1])
            return masks, labels
        elif self.aleatoric_head is not None:
            new_masks = self.aleatoric_head(decoder_output)
            return masks, new_masks

        return masks

[hyoseok1223]

MMSegmentation 에서 새로운 loss를 정의해서 사용할 때 마찬가지로 forward기반의 방법을 사용하는데, smp에 적용하시는게 불편하시다면 mm에서 loss를 추가하신 후에 사용해보시는 것도 괜찮을 것 같습니다.

MMSegmentatoin New Loss