src/slimface/training/accelerate_train.py

import os
import sys
import torch
import torch.nn as nn
import torchvision
import torch.multiprocessing as mp
import torchvision.transforms as transforms
from torch.utils.data import Dataset, DataLoader
from PIL import Image
import json
import argparse
import warnings
import pytorch_lightning as pl
from pytorch_lightning import Trainer
from pytorch_lightning.callbacks import ModelCheckpoint, TQDMProgressBar
from tqdm import tqdm
import math
from torch.optim.lr_scheduler import LambdaLR
import torchvision.models as models
import yaml

# Append the parent directory's 'models/edgeface' folder to the system path
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))

from data.process_face import extract_and_save_faces
from models.classification_models.alls import FaceClassifier
# from models.detection_models import align

# os.chdir(os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..', '..')))

# # Function to resolve string paths to Python objects
def resolve_path(path):
    """Convert a string like 'module.submodule.function' to a Python callable object."""
    try:
        module_name, obj_name = path.rsplit('.', 1)
        module = __import__("torchvision." + module_name, fromlist=[obj_name])
        return getattr(module, obj_name)
    except Exception as e:
        raise ValueError(f"Failed to resolve path {path}: {e}")


# Load MODEL_CONFIGS from YAML file
def load_model_configs(yaml_path):
    try:
        with open(yaml_path, 'r') as file:
            config = yaml.safe_load(file)
        if 'models' in config:
            config = config['models']
        model_configs = {}
        for model_name, params in config.items():
            model_configs[model_name] = {
                'resolution': params['resolution'],
                'model_fn': resolve_path(params['model_fn']),
                'weights': params['weights'].split(".")[-1]
            }
        return model_configs
    except FileNotFoundError:
        raise FileNotFoundError(f"Configuration file {yaml_path} not found.")
    except Exception as e:
        raise ValueError(f"Error loading YAML configuration: {e}")

# def extract_and_save_faces(input_dir, output_dir, algorithm='yolo', resolution=224):
#     """Preprocess images using face alignment and cache them with specified resolution."""
#     if align is None:
#         raise ImportError("face_alignment package is required for preprocessing.")
#     os.makedirs(output_dir, exist_ok=True)
#     with warnings.catch_warnings():
#         warnings.filterwarnings("ignore", category=FutureWarning, message=".*rcond.*")
#         for person in sorted(os.listdir(input_dir)):
#             person_path = os.path.join(input_dir, person)
#             if not os.path.isdir(person_path):
#                 continue
#             output_person_path = os.path.join(output_dir, person)
#             os.makedirs(output_person_path, exist_ok=True)
#             skipped_count = 0
#             for img_name in tqdm(os.listdir(person_path), desc=f"Processing {person}"):
#                 if not img_name.endswith(('.jpg', '.jpeg', '.png')):
#                     continue
#                 img_path = os.path.join(person_path, img_name)
#                 output_img_path = os.path.join(output_person_path, img_name)
#                 if os.path.exists(output_img_path):
#                     skipped_count += 1
#                     continue
#                 try:
#                     aligned_result = align.get_aligned_face([img_path], algorithm=algorithm)
#                     aligned_image = aligned_result[0][1] if aligned_result and len(aligned_result) > 0 else None
#                     if aligned_image is None:
#                         print(f"Face detection failed for {img_path}, using resized original image")
#                         aligned_image = Image.open(img_path).convert('RGB')
#                     aligned_image = aligned_image.resize((resolution, resolution), Image.Resampling.LANCZOS)
#                     aligned_image.save(output_img_path, quality=100)
#                 except Exception as e:
#                     print(f"Error processing {img_path}: {e}")
#                     aligned_image = Image.open(img_path).convert('RGB')
#                     aligned_image = aligned_image.resize((resolution, resolution), Image.Resampling.LANCZOS)
#                     aligned_image.save(output_img_path, quality=100)
#             if skipped_count > 0:
#                 print(f"Skipped {skipped_count} images for {person} that were already processed.")

class FaceDataset(Dataset):
    """Dataset for loading pre-aligned face images."""
    def __init__(self, root_dir, transform=None, resolution=224):
        self.root_dir = root_dir
        self.transform = transform
        self.resolution = resolution
        self.image_paths = []
        self.labels = []
        self.class_to_idx = {}
        for idx, person in enumerate(sorted(os.listdir(root_dir))):
            person_path = os.path.join(root_dir, person)
            if os.path.isdir(person_path):
                self.class_to_idx[person] = idx
                for img_name in os.listdir(person_path):
                    if img_name.endswith(('.jpg', '.jpeg', '.png')):
                        self.image_paths.append(os.path.join(person_path, img_name))
                        self.labels.append(idx)

    def __len__(self):
        return len(self.image_paths)

    def __getitem__(self, idx):
        img_path = self.image_paths[idx]
        label = self.labels[idx]
        try:
            image = Image.open(img_path).convert('RGB')
            image = image.resize((self.resolution, self.resolution), Image.Resampling.LANCZOS)
        except Exception as e:
            print(f"Error loading {img_path}: {e}")
            image = Image.new('RGB', (self.resolution, self.resolution))
        if self.transform:
            image = self.transform(image)
        return image, label

class FaceClassifierLightning(pl.LightningModule):
    """PyTorch Lightning module for face classification."""
    def __init__(self, base_model, num_classes, learning_rate,
                 warmup_steps=1000, total_steps=100000, max_lr_factor=10.0,
                 model_name='efficientnet_b0'):
        super(FaceClassifierLightning, self).__init__()
        self.model = FaceClassifier(base_model, num_classes, model_name, MODEL_CONFIGS)
        self.criterion = nn.CrossEntropyLoss()
        self.learning_rate = learning_rate
        self.warmup_steps = warmup_steps
        self.total_steps = total_steps
        self.max_lr = learning_rate * max_lr_factor
        self.min_lr = 1e-6
        self.model_name = model_name
        self.save_hyperparameters("num_classes", "learning_rate", "warmup_steps", "total_steps", "max_lr_factor", "model_name")

    def forward(self, x):
        return self.model(x)

    def training_step(self, batch, batch_idx):
        images, labels = batch
        outputs = self(images)
        loss = self.criterion(outputs, labels)
        self.log('train_loss', loss, prog_bar=True, on_step=True, on_epoch=True, sync_dist=True)
        _, predicted = torch.max(outputs, 1)
        acc = (predicted == labels).float().mean()
        self.log('train_acc', acc, prog_bar=True, on_step=True, on_epoch=True, sync_dist=True)
        return loss

    def validation_step(self, batch, batch_idx):
        images, labels = batch
        outputs = self(images)
        loss = self.criterion(outputs, labels)
        self.log('val_loss', loss, prog_bar=True, on_step=True, on_epoch=True, sync_dist=True)
        _, predicted = torch.max(outputs, 1)
        acc = (predicted == labels).float().mean()
        self.log('val_acc', acc, prog_bar=True, on_step=True, on_epoch=True, sync_dist=True)
        return loss

    def on_validation_epoch_end(self):
        metrics = self.trainer.logged_metrics
        train_loss = metrics.get('train_loss_epoch', 0.0)
        train_acc = metrics.get('train_acc_epoch', 0.0)
        val_loss = metrics.get('val_loss_epoch', 0.0)
        val_acc = metrics.get('val_acc_epoch', 0.0)
        current_lr = self.optimizers().param_groups[0]['lr']
        print(f"\nEpoch {self.current_epoch + 1}: "
              f"Train loss: {train_loss:.4f}, Train acc: {train_acc:.4f}, "
              f"Val loss: {val_loss:.4f}, Val acc: {val_acc:.4f}, "
              f"Learning rate: {current_lr:.6e}")

    def configure_optimizers(self):
        optimizer = torch.optim.Adam(self.model.conv_head.parameters(), lr=self.learning_rate)
        def lr_lambda(step):
            if step < self.warmup_steps:
                return (self.max_lr - self.learning_rate) / self.warmup_steps * step + self.learning_rate
            progress = (step - self.warmup_steps) / float(max(1, self.total_steps - self.warmup_steps))
            cosine_lr = 0.5 * (1.0 + math.cos(math.pi * progress))
            lr = self.min_lr + (self.max_lr - self.min_lr) * cosine_lr
            return max(lr, self.min_lr) / self.learning_rate
        scheduler = LambdaLR(optimizer, lr_lambda=lr_lambda)
        return {
            "optimizer": optimizer,
            "lr_scheduler": {
                "scheduler": scheduler,
                "interval": "step",
                "frequency": 1,
            }
        }

    def save_full_model(self, save_path):
        """Save the full model (base_model + conv_head) in TorchScript format."""
        os.makedirs(os.path.dirname(save_path), exist_ok=True)
        scripted_model = torch.jit.script(self.model)
        torch.jit.save(scripted_model, save_path)
        print(f"Full model saved in TorchScript format to {save_path}")

    def save_classifier_head(self, save_path):
        """Save only the classifier head (conv_head)."""
        os.makedirs(os.path.dirname(save_path), exist_ok=True)
        torch.save(self.model.conv_head.state_dict(), save_path)
        print(f"Classifier head saved to {save_path}")

class CustomModelCheckpoint(ModelCheckpoint):
    def format_checkpoint_name(self, metrics, ver=None):
        metrics['epoch'] = metrics.get('epoch', 0) + 1
        return super().format_checkpoint_name(metrics, ver)

class CustomTQDMProgressBar(TQDMProgressBar):
    def get_metrics(self, trainer, pl_module):
        items = super().get_metrics(trainer, pl_module)
        items["epoch"] = trainer.current_epoch + 1
        return items
    def init_train_tqdm(self):
        bar = super().init_train_tqdm()
        bar.set_description(f"Training Epoch {self.trainer.current_epoch + 1}")
        return bar   
    def on_train_epoch_start(self, trainer, pl_module):
        super().on_train_epoch_start(trainer, pl_module)
        if self.train_progress_bar:
            self.train_progress_bar.set_description(f"Training Epoch {trainer.current_epoch + 1}")

def main(args):
    mp.set_start_method('spawn', force=True)
    
    # Load model configurations using the provided config_path
    global MODEL_CONFIGS
    MODEL_CONFIGS = load_model_configs(args.image_classification_models_config_path)
    
    # Get the resolution for the selected model
    if args.classification_model_name not in MODEL_CONFIGS:
        raise ValueError(f"Model {args.classification_model_name} not supported. Choose from {list(MODEL_CONFIGS.keys())}")
    resolution = MODEL_CONFIGS[args.classification_model_name]['resolution']
    
    train_cache_dir = os.path.join(args.dataset_dir, f"train_data_aligned_{args.classification_model_name}")
    val_cache_dir = os.path.join(args.dataset_dir, f"val_data_aligned_{args.classification_model_name}")
    print(f"Preprocessing training dataset with resolution {resolution}...")
    extract_and_save_faces(
        input_dir=os.path.join(args.dataset_dir, "train_data"),
        output_dir=train_cache_dir,
        algorithm=args.algorithm,
        resolution=resolution
    )
    print(f"Preprocessing validation dataset with resolution {resolution}...")
    extract_and_save_faces(
        input_dir=os.path.join(args.dataset_dir, "val_data"),
        output_dir=val_cache_dir,
        algorithm=args.algorithm,
        resolution=resolution
    )
    transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
    ])
    train_dataset = FaceDataset(root_dir=train_cache_dir, transform=transform, resolution=resolution)
    val_dataset = FaceDataset(root_dir=val_cache_dir, transform=transform, resolution=resolution)
    if len(train_dataset) == 0 or len(val_dataset) == 0:
        raise ValueError("Dataset is empty. Check dataset directory or preprocessing.")
    train_loader = DataLoader(
        train_dataset,
        batch_size=args.batch_size,
        shuffle=True,
        drop_last=True,
        num_workers=2,
        pin_memory=True,
        persistent_workers=True
    )
    val_loader = DataLoader(
        val_dataset,
        batch_size=args.batch_size,
        shuffle=False,
        drop_last=True,
        num_workers=2,
        pin_memory=True,
        persistent_workers=True
    )
    
    steps_per_epoch = len(train_loader)
    if steps_per_epoch == 0:
        raise ValueError("Train DataLoader is empty. Check dataset size or batch configuration.")
    total_steps = args.num_epochs * steps_per_epoch
    warmup_steps = int(args.warmup_steps * total_steps) if args.warmup_steps > 0 else int(0.05 * total_steps)
    
    # Load the appropriate model
    model_fn = MODEL_CONFIGS[args.classification_model_name]['model_fn']
    weights = MODEL_CONFIGS[args.classification_model_name]['weights']
    base_model = model_fn(weights=weights)
    
    for param in base_model.parameters():
        param.requires_grad = False
    if hasattr(base_model, 'classifier'):
        base_model.classifier = nn.Identity()
    elif hasattr(base_model, 'fc'):
        base_model.fc = nn.Identity()
    elif hasattr(base_model, 'head'):
        base_model.head = nn.Identity()
    base_model.eval()
    
    model = FaceClassifierLightning(
        base_model=base_model,
        num_classes=len(train_dataset.class_to_idx),
        learning_rate=args.learning_rate,
        warmup_steps=warmup_steps,
        total_steps=total_steps,
        max_lr_factor=args.max_lr_factor,
        model_name=args.classification_model_name
    )

    ckpts_backup_dir = './ckpts/ckpts_backup'
    os.makedirs(ckpts_backup_dir, exist_ok=True)
    checkpoint_callback = CustomModelCheckpoint(
        monitor='val_loss',
        dirpath=ckpts_backup_dir,
        filename=f'SlimFace_{args.classification_model_name}_{{epoch:02d}}_{{val_loss:.2f}}',
        save_top_k=1,
        mode='min'
    )
    
    progress_bar = CustomTQDMProgressBar()
    trainer = Trainer(
        max_epochs=args.num_epochs,
        accelerator=args.accelerator,
        devices=args.devices,
        callbacks=[checkpoint_callback, progress_bar],
        log_every_n_steps=10
    )
    trainer.fit(model, train_loader, val_loader)
    
    # Save the idx_to_class mapping
    idx_to_class = {v: k for k, v in train_dataset.class_to_idx.items()}
    idx_to_class_path = os.path.join('./ckpts', 'index_to_class_mapping.json')
    os.makedirs(os.path.dirname(idx_to_class_path), exist_ok=True)
    with open(idx_to_class_path, 'w') as f:
        json.dump(idx_to_class, f, indent=4)
    print(f"Index to class mapping saved to {idx_to_class_path}")

    # Save the full model and classifier head after training
    full_model_save_path = os.path.join('./ckpts', f'SlimFace_{args.classification_model_name}_full_model.pth')
    classifier_head_save_path = os.path.join('./ckpts', f'SlimFace_{args.classification_model_name}_conv_head.pth')
    
    model.save_full_model(full_model_save_path)
    # model.save_classifier_head(classifier_head_save_path)

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Train a face classification model with PyTorch Lightning.')
    parser.add_argument('--dataset_dir', type=str, default='./data/processed_ds',
                        help='Path to the dataset directory.')
    parser.add_argument('--image_classification_models_config_path', type=str, default='./configs/image_classification_models_config.yaml',
                        help='Path to the YAML configuration file for model configurations.')
    parser.add_argument('--batch_size', type=int, default=8,
                        help='Batch size for training and validation.')
    parser.add_argument('--num_epochs', type=int, default=100,
                        help='Number of training epochs.')
    parser.add_argument('--learning_rate', type=float, default=5e-4,
                        help='Initial learning rate for the optimizer.')
    parser.add_argument('--max_lr_factor', type=float, default=10.0,
                        help='Factor to multiply initial learning rate to get max learning rate during warmup.')
    parser.add_argument('--accelerator', type=str, default='auto',
                        choices=['cpu', 'gpu', 'tpu', 'auto'],
                        help='Accelerator type for training.')
    parser.add_argument('--devices', type=int, default=1,
                        help='Number of devices to use (e.g., number of GPUs).')
    parser.add_argument('--algorithm', type=str, default='yolo',
                        choices=['mtcnn', 'yolo'],
                        help='Face detection algorithm to use (mtcnn or yolo).')
    parser.add_argument('--warmup_steps', type=float, default=0.05,
                        help='Fraction of total steps for warmup phase (e.g., 0.05 for 5%).')
    parser.add_argument('--total_steps', type=int, default=0,
                        help='Total number of training steps (0 to use epochs * steps_per_epoch).')
    parser.add_argument('--classification_model_name', type=str, default='efficientnet_b0',
                        choices=list(load_model_configs('./configs/image_classification_models_config.yaml').keys()),
                        help='Model to use for training.')

    args = parser.parse_args()
    main(args)