pytorch_pretrained_model.py

learning_rate = 0.0001

#get mobilenet
model = torchvision.models.mobilenet_v2(pretrained=True)

# freeze the feature extraction convolutional layers
for param in model.parameters():
    param.requires_grad = False

# define a classification layer
model.classifier[1] = nn.Linear(in_features=model.classifier[1].in_features, out_features=1)

# send the model to the device
model = model.to(device)

# only optimize the classification layer
optimizer = torch.optim.Adam(model.classifier[1].parameters(), lr=learning_rate)