arcface.py 391 Bytes
Newer Older
Tiago de Freitas Pereira's avatar
Update  
Tiago de Freitas Pereira committed
1
2
3
4
5
6
7
8
9
10
11
from bob.bio.face.pytorch.head import ArcFace


### train_dataloader AND BACKBONES needs to be loaded before
#NUM_CLASS = train_dataloader.dataset.n_classes
#WEIGHT = train_dataloader.dataset.get_demographic_class_weights()
NUM_CLASS = len(list(train_dataloader.dataset.demographic_keys.values()))

identity_head = ArcFace(feat_dim=backbone.features.num_features, num_class=NUM_CLASS)

pass