Retrain facenet neural network

user_5404647 0
Viewed 73

I want to use the pretained MTCNN model to train in a subset of images, and keep the layers that are important to get the features of the images. Also, I want to use another loss function.

import torch
from facenet_pytorch import InceptionResnetV1, MTCNN
from torch.utils.data import DataLoader
from torchvision import datasets
import numpy as np
import pandas as pd
import os

workers = 0 if os.name == 'nt' else 4
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
print('Running on device: {}'.format(device))

mtcnn = MTCNN(
    image_size=160, margin=0, min_face_size=20,
    thresholds=[0.6, 0.7, 0.7], factor=0.709, post_process=True,
    device=device
)

The problem is that I don't know how to face the problem of "freezing" some layers of the neural network that I don't want to tune and tune only the last layers. Also, how to apply the new criterion to this pretrained network.

pytorch python
2 Answers

One idea is to stop gradient calculation of the mtcnn model ...

for param in mtcnn.parameters():
    param.requires_grad = False

You can then replace the final layer (or as many you need) of mtcnn with ones that you create ...

mtcnn.fc = nn.Linear(num_in_features, my_num_out_features)

Then train as usual, only fc layer will be trained.

user_9671620

You can first print out or visualize the MTCNN model, so you will know the name (key) of the layer you want to freeze. Assuming the name of the layer is feature_layer, you can then simple do:

for name, param in model.named_parameters():
    if not 'feature_layer' in name:
        param.requires_grad = False

So it will freeze every other layers except the last layer feature_layer. Then your other parameters will follow suit:

optimizer = torch.optim.Adam(params, lr=1e-3)
lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.1)

Now, to change the criterion, you change it here to one of the default ones, or use a custom one:

loss_fn = torch.nn.CrossEntropyLoss()
user_18360690
Related