arcface.py

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