[FIXED] Retrain facenet neural network

Issue

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.

Solution

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()

Answered By - Olasimbo Arigbabu