Vision: Binary Classification on Resnet?

[Linardos]

I want to classify 2 classes of images from ImageNet so I do the following:

  model = models.__dict__['resnet18'](pretrained=False, num_classes=2)

However this gives an output of dimensions (batch_size, 2). When dealing with 2 classes I would like to have a 1 dimensional vector output instead, so how would I have resnet do that? (I tried changing the num_classes to 1 but got an error).

[Linardos]

Turns out it works if you change the number of classes to 1 but you have to make sure you use
BCEWithLogitsLoss instead of CrossEntropyLoss and then unsqueeze and transpose the target vector so that they have the same shape

[alykhantejani]

Closing as seems to be resolve. @Linardos please comment back on this issue if you have any further questions.

[oya163]

I believe yes you can use CrossEntropyLoss. I followed this tutorial and its pretty easy.
http://pytorch.org/tutorials/beginner/transfer_learning_tutorial.html
It is classifying bees and ants images.

[fermat97]

I have modified the resnet18 for binary classification as:
model = torchvision.models.resnet18() model.conv1 = nn.Conv2d(num_input_channel, 64, kernel_size=7, stride=2, padding=3,bias=False) model.avgpool = nn.AdaptiveAvgPool2d(1) model.fc = nn.Linear(512 * torchvision.models.resnet.BasicBlock.expansion,2)

and I have tried this with nn.CrossEntropyLoss(). I think there is a problem with this classifier due to its poor performance (it is even worse than a dummy classifier). @Linardos I couldn't understand what is the problem with the loss and number of classes?