How to use mobilnet v2 for binary classification in pytorch?

Question

I'm trying to train a binary classifier using transfer learning in mobilenet v2 but am not sure how to freeze the layers and make it classify between 0 and 1. Any help would be appreciated

Answer 1

To use pretrained model with 2 outputs:

1. torchvision

import torch
from torchvision.models import mobilenet_v2

model = mobilenet_v2(pretrained=True)
model.classifier[1] = torch.nn.Linear(model.classifier[1].in_features, 2)

2. timm:

import timm
model = timm.create_model('mobilenetv2_100', pretrained=True, num_classes=2)

To freeze all parameters:

model.requires_grad_(False)

To freeze specific parameters:

import torch
from torchvision.models import mobilenet_v2

model = mobilenet_v2(pretrained=True)
model.classifier[1] = torch.nn.Linear(model.classifier[1].in_features, 2)

for name, param in model.named_parameters():
    if "classifier" in name:
        param.requires_grad = True
    else:
        param.requires_grad = False

Here we freeze all but the last fully connected layer that can be fine-tuned for binary classification.

Full example with training loop:

import torch
from torchvision.models import mobilenet_v2

model = mobilenet_v2(pretrained=True)
model.classifier[1] = torch.nn.Linear(model.classifier[1].in_features, 2)

for name, param in model.named_parameters():
    if "classifier" in name:
        param.requires_grad = True
    else:
        param.requires_grad = False


torch.set_grad_enabled(True)
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
loss_fn = torch.nn.CrossEntropyLoss()
x = torch.rand(size=(4, 3, 224, 224))
y = torch.randint(0, 2, size=(4,))

for _ in range(10):
    model.zero_grad()
    y_pred = model(x)
    loss = loss_fn(y_pred, y)
    loss.backward()
    optimizer.step()
    print(loss.item())