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 != "onet.dense6_3.bias":
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.