Swin-Transformer icon indicating copy to clipboard operation
Swin-Transformer copied to clipboard

Is multi-label classification for image classification supported?

Open Terry-Kusunoki-Martin opened this issue 3 years ago • 2 comments
trafficstars

I am trying to train an image classifier where image ground truth contains multiple classes. Is it possible to train a model that outputs multiple classes?

Terry-Kusunoki-Martin avatar Sep 07 '22 14:09 Terry-Kusunoki-Martin

I am trying to train an image classifier where image ground truth contains multiple classes. Is it possible to train a model that outputs multiple classes?

Yes, it should be not hard to adapt the code to support multiple labels. You need to modify the origial loss to either sigmoid binary cross-entropy loss or softmax based soft cross-entropy loss.

ancientmooner avatar Sep 29 '22 15:09 ancientmooner

@ancientmooner appreciate the reply!

So I would be modifying this block of code here to add a BCE loss function? https://github.com/microsoft/Swin-Transformer/blob/afeb877fba1139dfbc186276983af2abb02c2196/main.py#L109

I also would need to change the output layer to match the number of classes right? Where would I make that modification?

Terry-Kusunoki-Martin avatar Oct 07 '22 13:10 Terry-Kusunoki-Martin