How to completely modify the classifier of a pretrained network to use KNN

1

Hi, I’m currently using a pretrained MobileNetV2 model for my audio classification.
As a beginner of using PyTorch and from reading other similar topics, I understand how to modify the first and the last layer:

from torchvision.models import mobilenet_v2
import torch
import torch.nn as nn
mobilenet_model = mobilenet_v2(pretrained=True)
mobilenet_model.features[0][0] = nn.Conv2d(1, 32, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)  # Modify the first layer
mobilenet_model.classifier[1] = nn.Linear(1280, NUM_CLASSES, bias=True) # Modify the last layer

It works great, but now I want to experiment this model with different approaches by applying different classifiers (like KNN), but I’m not sure how to change the last layer to use, let’s say, KNN instead of linear transformation.

Does anyone know how to do this? Any resource is also fine!

2

If you want to freeze the model and use a KNN implementation from e.g. sklearn, you could replace the last linear layer with nn.Identity to get the features from the penultimate layer and feed them to sklearn.neighbors.KNeightborsClassifier.
Alternatively, you could also implement KNN manually in PyTorch (you should be able to find some implementations in this forum or probably on GitHub).

1 Like