Skip to content

tungedng2710/Improved-ArcFace

Repository files navigation

Improved ArcFace: Some improvements on ArcFace model

Reference: https://github.com/deepinsight/insightface (Official repository of ArcFace)

Preparation

Create your dataset like the sample dataset:

data
--YourDatasetName
  --Label 1
  ----image1.jpg
  ----image2.jpg
  ----...
  --Label 2
  ----...

You can find pretrained model at link

For face alignment, run

python align_face.py --root_dir /path/to/dataset/folder --dst_w 112 --dst_h 112

(You can change the destination size --dst_w is output width and --dst_h is output height)

Configuration

Train

  • loss: Now you can choose ArcFace or ElasticArcFace.
  • backbone: Find supported backbone in ArcFaceModel's docstring. irse50 and mobilefacenet have pretrained models on insightface's datasets at link. Other ones are listed in the doctring of ArcFaceModel class, but they have to be trained from scratch.
  • freeze_model: Freeze the backbone of trained model for transfer learning. Set its value is false to train from scratch.
  • root_dir: The path to the directory of train dataset
  • use_lr_scheduler: Use learning rate scheduler.
  • optimizer: Optimizer for training progress. sam, lamb, adam and adan are supported optimizer. You can find the original implementation of SAM, LAMB and Adan. If you don't change the config, the default optimizer is ADAM.
  • verbose: 0: nothing will be shown; 1: shows results per epoch only; 2: shows train losses per iteration
  • prefix: Prefix of saved model's name.

Test

  • trainset_path: Path to the directory of dataset used for training the test model.
  • testset_path: Path to the directory of test dataset.
  • pretrained_model_path: Path to the pretrained model.

Training

In terminal, run

python main.py --config ./path/to/config/file.json --phase train --device [gpu_id]

Currently, only training on a single gpu is supported.

Testing

In terminal, run

python main.py --config ./path/to/config/file.json --phase test --device [gpu_id]

Verification

In terminal, run

python verification.py --config ./path/to/config/file.json

Inference on single image is also supported, check the commented lines in the verification.py file. Filling hair is a miscellaneous option to observe the decrease of model's performance when the given face is partly covered.