Retraining pytorch model (augmented learning)
I have the following code:
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
)

def collate_fn(x):
    return x[0]

dataset = datasets.ImageFolder('data/images/')
dataset.idx_to_class = {i:c for c, i in dataset.class_to_idx.items()}
loader = DataLoader(dataset, collate_fn=collate_fn, num_workers=workers)
#print(dataset.idx_to_class)

aligned = []
names = []
i = 0
for x, y in loader:
    x_aligned, prob = mtcnn(x, return_prob=True)
    if x_aligned is not None:
        print('Face detected with probability: {:8f}'.format(prob))
        aligned.append(x_aligned)
        names.append(dataset.idx_to_class[y])
        i += 1
#print(i)

for name, param in mtcnn.named_parameters(): #Freezing everything but last layer
    #print(name)
    if name != &quot;onet.dense6_3.bias&quot;:
        param.require_grad = False
    else:
        param.require_grad = True

And now I would like to retrain this model to predict three classes (Now it only predicts the probability of a face). Let say that I have inside data/images/ three folders, faces1, faces2 and faces3. How could I retrain this model with these three folders? I would like to have a tensor like [prob1, prob2, prob3] with the probability of an image for each class. Thanks.

MTCN: This class loads pretrained P-, R-, and O-nets and returns images cropped to include the face only, given raw input images.
I am assuming that you are trying to use InceptionResnetV1 for classification on your dataset. To retrain the Inception model you just load the model with the number of classes you need and then train it.
resnet = InceptionResnetV1(
    classify=True,
    pretrained='vggface2',
    num_classes=3
)

Complete finetuning example is here https://github.com/timesler/facenet-pytorch/blob/master/examples/finetune.ipynb
