Sigmoid, Softmax in semantic segmentation

semantic segmentation을 공부하면서 최종 output feature map에 어떤 activation function을 취해야 하는지 헷갈렸던 내용을 정리합니다.

 

 

 

• final output feature map의 사이즈는 (C, H, W)이고 아래와 같은 값이 나왔다고 가정하겠습니다.
• 여기선 dog, cat, person, backgorund 4가지 클래스가 있고 각 채널은 특정 클래스에 대응됩니다.

 

 

 

 

 

• Sigmoid를 취하게 되면 각 element마다 0~1사이의 값으로 출력됩니다. element 단위로 계산되기 때문에 각 픽셀에 대한 독립적인 확률값을 구할 수 있습니다.

 

 

 

 

 

• 하나의 이미지에 하나의 라벨만 존재하는 image classification 문제의 경우 softmax를 취하고 가장 큰 값을 예측 클래스로 사용합니다.
• semantic segmentation 또한 분류 단위가 이미지에서 픽셀 단위로 바뀌었을 뿐 개념은 같습니다. 이런 상황에서  (0, 1)에 위치한 픽셀은 cat으로 봐야하는지, background로 봐야하는지 애매합니다.
• 여기선 background의 값이 더 크기 때문에 background로 판단할 수 있겠지만 두 픽셀 값의 차이가 거의 없는 경우 어떤 클래스로 판단해야 하는지 애매합니다.
• 일반적인 semantic segmentation task에서 하나의 픽셀은 하나의 클래스에만 대응되기 때문에 sigmoid를 사용하게 되면 이러한 문제에 직면할 수 있습니다. (도메인에 따라 semantic segmentation task에서 하나의 픽셀이 여러 라벨을 갖는 multi-label classification 문제도 존재합니다.)

 

 

 

 

 

• 이번엔 sigmoid대신 softmax를 취하는 경우를 보겠습니다.
• pytorch에선 softmax를 취할 때 어떤 축을 기준으로 계산할건지 파라미터를 줄 수 있습니다.
• 만약 채널 축을 기준으로 softmax를 취하게 되면 같은 위치에 있으면서 다른 채널에 존재하는 element들의 합이 1이 되도록 계산합니다. 

 

 

 

 

 

• 모든 채널에서 같은 위치에 있는 값들을 더했을 때 1이 됩니다.

 

 

 

 

 

 

 

 

 

 

• 최종 예측에선 각 채널 중 가장 큰 값이 해당 픽셀의 클래스가 됩니다.
• (0, 0)에서 가장 큰 확률 값은 cat 클래스에 해당하므로 (0, 0)은 cat으로 분류하게 됩니다.

 

 

 

 

 

 

• 위 예시에선 채널 축을 기준으로 softmax를 취했는데 height로 취하게 되면 어떻게 계산되는지 보겠습니다.
• height를 기준으로 하면 각 채널마다 독립적으로 계산되며, matrix의 각 column마다 합이 1이 되도록 합니다.

 

 

 

 

 

 

 

 

• 이번에는 width 축을 기준으로 softmax를 취하는 경우입니다.
• 채널마다 독립적으로 계산되고, matrix의 각 row마다 합이 1이 되도록 계산합니다.

 

 

 

 

 

 

 

 

• 아래 내용은 SegNet의 section 3의 내용입니다.
• 마지막 문장에서 픽셀마다 독립적인 확률값을 계산하기 위해 softmax를 취한다고 했는데 그러면 sigmoid를 취했을 때 결과랑 똑같은 거 아닌가? 라는 생각을 했습니다.
• 제가 해당 내용을 잘못 받아들였던 이유는 feature map이 softmax, sigmoid를 거쳤을 때 어떤 결과가 나오는지 그 차이점을 잘못 이해하고 있었기 때문입니다.

 

 

 

 

 

• 픽셀 단위로 독립적인 softmax를 취했을 때, 아래와 같이 각 픽셀마다 0~1 사이의 확률 값이 나온다고 생각했습니다.
• sigmoid를 사용해도 각 픽셀마다 0~1 사이의 확률 값이 나오기 때문에 결과적으로 차이가 없다고 생각했습니다.

 

 

 

 

• 그러나 각 픽셀마다 softmax를 취하게 되면 하나의 원소만 존재하므로 출력값은 항상 1이 됩니다.

 

 

 

 

 

 

 

• 따라서, k개의 클래스가 존재하는 semantic segmentation task에선 (K, H, W) 최종 output feature map에 채널 축을 기준으로 softmax를 사용해야합니다.

 

 

 

 

 

 

 

 

 

 

 

 

 

 

 

 

 

 

 

• 그런데 pytorch를 사용한 구현체들을 보면 softmax를 사용하지 않습니다.
• softmax를 사용해야 하는 것이 맞는데, 왜 대다수의 pytorch 구현체들은 softmax layer가 존재하지 않는지 이해할 수 없었습니다.

https://github.com/usuyama/pytorch-unet/blob/master/pytorch_unet.py

 

 

 

• 이 부분에 대한 갈증은 pytorch document에서 crossentropyloss를 확인하고서야 해소할 수 있었습니다.
• pytorch에선 cross entropy loss에 softmax 결과를 포함하고 있기 때문에 모델 구현 단계에서 굳이 softmax layer를 사용하지 않아도 됩니다.