transfer learning or fine tuning insightface on custom (my own) dataset

Author: acel122Created Mar 18, 2025Updated May 14, 2026

I'm currently working on transfer learning with InsightFace using the glint360k_cosface_r100_fp16_0.1 model from the ArcFace Torch section. However, I'm facing issues with either overfitting or underfitting on my dataset, and I'm not sure what I'm doing wrong. Here are the problems I'm encountering:

My dataset consists of 127 individuals, with only 7 images per person taken from different angles: front view, 3/4 left and right, upper view, lower view, and profile left and right. This results in a total dataset of 889 images.
Initially, I split the data 80% for training and 20% for validation at the folder level, meaning each person has 5 training images and 2 validation images. This led to underfitting, likely due to insufficient training data per person.
To address this, I performed data augmentation first and then applied the same 80/20 split. However, this resulted in overfitting, as I suspect the model was "cheating" by memorizing patterns from the augmented images rather than generalizing. Below is the pseudocode representing my approach:

BEGIN

# ---- SETUP ENVIRONMENT ----
SET CUDA and OpenCV paths
SET PyTorch memory allocation config

# ---- IMPORT LIBRARIES ----
IMPORT required libraries (Torch, NumPy, OpenCV, InsightFace, etc.)

# ---- DEFINE FaceDataset CLASS ----
CLASS FaceDataset:
    INITIALIZE dataset directory, transformations, and cache
    IF cache exists:
        LOAD dataset from cache
    ELSE:
        INITIALIZE face detection model (InsightFace)
        SCAN dataset directory
        FOR each image folder:
            FOR each image:
                DETECT face
                IF face detected:
                    CROP and RESIZE to (112,112)
                    STORE in dataset
        SAVE dataset to cache
    
    FUNCTION _detect_face(image):
        READ image
        CONVERT to RGB
        DETECT faces using InsightFace
        IF face detected:
            CROP, RESIZE, RETURN face
        ELSE:
            RETURN None

    FUNCTION __getitem__(index):
        RETURN image and label

    FUNCTION __len__():
        RETURN number of samples

# ---- DEFINE FaceRecognitionModel CLASS ----
CLASS FaceRecognitionModel:
    INITIALIZE ResNet50 backbone
    FREEZE lower layers, fine-tune upper layers
    ADD fully connected classifier with dropout
    FUNCTION forward(input):
        PASS through backbone
        PASS through classifier head
        RETURN output

# ---- DEFINE TRAINING FUNCTION ----
FUNCTION train_model(model, train_loader, val_loader, criterion, optimizer, scheduler, num_epochs):
    INITIALIZE metrics storage
    SET early stopping threshold
    
    FOR epoch in range(num_epochs):
        IF warm-up phase:
            ADJUST learning rate
        
        # ---- TRAIN PHASE ----
        SET model to training mode
        FOR batch in train_loader:
            LOAD input images and labels
            COMPUTE predictions
            CALCULATE loss
            BACKPROPAGATE and update weights

        # ---- VALIDATION PHASE ----
        SET model to evaluation mode
        FOR batch in val_loader:
            COMPUTE predictions
            CALCULATE validation loss
        UPDATE scheduler with validation loss
        CHECK for early stopping condition

    RETURN best model

# ---- DEFINE VALIDATION ----
FUNCTION split_data():
    EXTRACT person identity from filenames
    PERFORM GroupShuffleSplit to avoid identity leakage
    RETURN train and validation indices

# ---- DEFINE DATA AUGMENTATION ----
FUNCTION get_transforms():
    RETURN image augmentation pipeline (flip, resize, normalize)

# ---- DEFINE ONNX EXPORT FUNCTION ----
FUNCTION export_to_onnx(model, save_path):
    CONVERT PyTorch model to ONNX format
    VERIFY conversion
    RETURN ONNX model

# ---- MAIN FUNCTION ----
FUNCTION main():
    SET dataset path, cache directory, and logging path
    INITIALIZE dataset with caching enabled
    SPLIT dataset ensuring unique individuals in train and validation
    APPLY data augmentation
    CREATE data loaders for training and validation
    
    # ---- MODEL INITIALIZATION ----
    LOAD ResNet50 backbone
    INITIALIZE FaceRecognitionModel
    SET loss function, optimizer, and scheduler
    
    # ---- TRAIN THE MODEL ----
    CALL train_model()
    
    # ---- EXPORT TRAINED MODEL ----
    CALL export_to_onnx()

    PRINT "Training Complete!"

# ---- RUN MAIN FUNCTION ----
IF __name__ == "__main__":
    CALL main()

END

Since I'm new to this field, I would really appreciate detailed feedback on what mistakes I might be making in my approach or code. Thanks in advance!

Source: deepinsight/insightface

View original on GitHub View discussion on GitHub