config.py

configurations = {
    1: dict(
        SEED=1337,
        SYNC_DATA=True,
        SYNC_DATA_NUMCLASS=100 * 10000,  # 2059904,
        DATA_ROOT="../data/webface260m",
        VAL_SET="lfw, cfp_fp, agedb_30",
        MODEL_ROOT="../models/tmp/model",
        IS_RESUME=False,
        BACKBONE_RESUME_ROOT="",
        HEAD_RESUME_ROOT="",
        BACKBONE_NAME="IR_100",  # support: ['MobileFaceNet', 'ResNet_50', 'ResNet_101', 'ResNet_152',
        #'IR_50', 'IR_101', 'IR_152', 'IR_SE_50', 'IR_SE_100', 'IR_SE_101', 'IR_SE_152',
        #'AttentionNet_IR_56', 'AttentionNet_IRSE_56','AttentionNet_IR_92', 'AttentionNet_IRSE_92',
        #'ResNeSt_50', 'ResNeSt_101', 'ResNeSt_100']
        HEAD_NAME="ParallelArcFace",  # ArcFace, ParallelArcFace
        LOSS_NAME="ParallelArcLoss",  # CrossEntropy, ParallelArcLoss
        ENABLE_MODEL_PARALLEL=True,
        MODEL_PARALLEL_DIM=1,  # -1/1/0
        ENABLE_AMP=True,
        OPTIMIZER="sgd",
        ENABLE_ZERO_OPTIM=True,
        INPUT_SIZE=[112, 112],
        EMBEDDING_SIZE=512,
        BATCH_SIZE=512,
        EVAL_FREQ=2000,
        LR=0.1,
        LR_SCHEDULER="cosine",
        WARMUP_EPOCH=0,
        WARMUP_LR=0.0,
        START_EPOCH=0,
        NUM_EPOCH=16,
        LR_STEP_SIZE=10,
        LR_DECAY_EPOCH=[10, 18, 22],
        LR_DECAT_GAMMA=0.1,
        LR_END=1e-5,
        WEIGHT_DECAY=5e-4,
        MOMENTUM=0.9,
        WORLD_SIZE=1,
        RANK=0,
        GPU=[0, 1, 2, 3, 4, 5, 6, 7],
        DIST_BACKEND="nccl",
        DIST_URL="tcp://localhost:23456",
        NUM_WORKERS=1,
        #
        #
        #
        # Data Augmentation
        RANDAUGMENT=False,
        RANDAUGMENT_N=2,
        RANDAUGMENT_M=9,
        RANDOM_ERASING=False,
        MIXUP=False,
        MIXUP_ALPHA=1.0,
        MIXUP_PROB=0.5,
        CUTOUT=False,
        CUTMIX=False,
        CUTMIX_ALPHA=1.0,
        CUTMIX_PROB=0.5,
        COLORJITTER=False,
    ),
}