Issue
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.
Solution
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
Answered By - Bhupen
0 comments:
Post a Comment
Note: Only a member of this blog may post a comment.