facial_attribute_classification_final.py

# -*- coding: utf-8 -*-
"""Facial-Attribute-Classification-Final.ipynb

Automatically generated by Colab.

Original file is located at
    https://colab.research.google.com/drive/1Mn2pvx_mTYJU5EdcS7QMUuMJcnc8N45J

## **Facial Attribute Classification using CNN + RNN with Attention**
**Neural Computing & Deep Learning Project**

---

### **Step 1: Install Required Libraries**
Before starting, we need to install the necessary libraries. These include:
- `transformers`: For image preprocessing using the ViT (Vision Transformer) processor.
- `datasets`: To load the CelebA dataset from Hugging Face.
- `torch` and `torchvision`: For building and training the deep learning model.
- `matplotlib`: For visualizing results.
- `scikit-learn`: For stratified train/validation splitting.
"""

"""!pip install transformers datasets torch torchvision matplotlib scikit-learn"""

"""### **Step 2: Import Libraries**
We import all the required libraries and modules. This includes:
- PyTorch for model building and training.
- Torchvision for image transformations and pretrained models.
- Hugging Face `datasets` and `transformers` for dataset loading and preprocessing.
- Matplotlib for visualization.
"""

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision import transforms, models
from torchvision.transforms import v2
from datasets import load_dataset
from transformers import AutoImageProcessor
import matplotlib.pyplot as plt
from sklearn.model_selection import train_test_split
from collections import defaultdict
from torch.optim.lr_scheduler import ReduceLROnPlateau

"""### **Step 3: Check for CUDA Availability**
We check if a GPU (CUDA) is available. If not, the code will fall back to the CPU. This ensures compatibility across different hardware setups.

"""

# Check for CUDA availability
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device}")

"""### **Step 4: Load and Prepare the Dataset**
We load the CelebA dataset from Hugging Face and prepare it for training and validation. Key steps include:
- Loading a subset of the dataset (10% for efficiency).
- Splitting the data into training and validation sets using stratified sampling to maintain class balance.
- Applying image transformations (augmentation for training, basic preprocessing for validation).

"""

def load_and_prepare_data():
    # Load 10% of the CelebA dataset
    full_data = load_dataset("flwrlabs/celeba", split="train[:1%]")

    # Stratified split to maintain class balance
    train_idx, val_idx = train_test_split(
        range(len(full_data)),
        test_size=0.2,
        stratify=full_data["Male"]  # Use the "Male" attribute for stratification
    )

    train_data = full_data.select(train_idx)
    val_data = full_data.select(val_idx)

    # Define image transformations
    train_transforms = v2.Compose([
        v2.RandomHorizontalFlip(p=0.5),  # Randomly flip images horizontally
        v2.ColorJitter(brightness=0.2, contrast=0.2),  # Adjust brightness and contrast
        v2.ToDtype(torch.float32, scale=True)  # Convert to tensor and normalize
    ])

    val_transforms = v2.Compose([
        v2.ToDtype(torch.float32, scale=True)  # Basic preprocessing for validation
    ])

    # Preprocessing function
    def preprocess(batch, augment=False):
        processor = AutoImageProcessor.from_pretrained("google/vit-base-patch16-224", use_fast=True)  # Fast image processor
        trans = train_transforms if augment else val_transforms
        images = [trans(image.convert("RGB")) for image in batch["image"]]  # Convert to RGB and apply transformations
        batch["pixel_values"] = [processor(img, return_tensors="pt")["pixel_values"] for img in images]
        batch["labels"] = batch["Male"]  # Use the "Male" attribute as the target label
        return batch

    return (
        train_data.with_transform(lambda x: preprocess(x, augment=True)),  # Augment training data
        val_data.with_transform(preprocess)  # No augmentation for validation
    )

train_data, val_data = load_and_prepare_data()

"""### **Step 5: Create Dataset and DataLoader**
We define a custom `CelebADataset` class to handle the dataset and create `DataLoader` objects for training and validation. This ensures efficient batching and shuffling of data during training.

"""

# Dataset class
class CelebADataset(torch.utils.data.Dataset):
    def __init__(self, dataset):
        self.dataset = dataset

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

    def __getitem__(self, idx):
        item = self.dataset[idx]
        return item["pixel_values"].squeeze(0), torch.tensor(item["labels"], dtype=torch.float32)

# Create DataLoaders
train_loader = DataLoader(CelebADataset(train_data), batch_size=32, shuffle=True, num_workers=2)
val_loader = DataLoader(CelebADataset(val_data), batch_size=32, num_workers=2)

"""### **Step 6: Define the Model Architecture**
We define a hybrid CNN-RNN model with attention. The model consists of:
- A **pretrained ResNet-18 backbone** for feature extraction.
- A **bidirectional LSTM** to capture spatial dependencies.
- A **multi-head attention mechanism** to focus on important regions.
- A **fully connected classifier** for binary classification.

"""

# Multi-head attention module
class MultiHeadAttention(nn.Module):
    def __init__(self, embed_dim, num_heads):
        super().__init__()
        self.attn = nn.MultiheadAttention(embed_dim, num_heads, batch_first=True)

    def forward(self, x):
        attn_output, _ = self.attn(x, x, x)
        return attn_output.mean(dim=1)

# Main model
class FaceRecognitionModel(nn.Module):
    def __init__(self):
        super().__init__()
        # Backbone (ResNet-18)
        resnet = models.resnet18(weights=models.ResNet18_Weights.IMAGENET1K_V1)
        self.backbone = nn.Sequential(*list(resnet.children())[:-3])  # Use up to layer3

        # Freeze early layers
        for param in self.backbone[:5].parameters():
            param.requires_grad = False

        # Adaptive pooling
        self.adaptive_pool = nn.AdaptiveAvgPool2d((14, 14))

        # RNN + Attention
        self.rnn = nn.LSTM(input_size=256, hidden_size=128, bidirectional=True, batch_first=True)
        self.attention = MultiHeadAttention(embed_dim=256, num_heads=4)
        self.classifier = nn.Sequential(
            nn.Linear(256, 64),
            nn.ReLU(),
            nn.Dropout(0.2),
            nn.Linear(64, 1)
        )

    def forward(self, x):
        # CNN features
        x = self.backbone(x)
        x = self.adaptive_pool(x)

        # Reshape for RNN
        batch_size = x.size(0)
        x = x.view(batch_size, -1, 256)  # (B, 14*14=196, 256)

        # Sequence processing
        rnn_out, _ = self.rnn(x)
        attn_out = self.attention(rnn_out)
        return self.classifier(attn_out)

# Initialize model
model = FaceRecognitionModel().to(device)

"""### **Step 7: Training Setup**
We define the loss function, optimizer, and learning rate scheduler. Mixed precision training is enabled if CUDA is available.

"""

# Loss function and optimizer
criterion = nn.BCEWithLogitsLoss()
optimizer = optim.AdamW(model.parameters(), lr=1e-4, weight_decay=1e-4)

# Mixed precision training (if CUDA is available)
if device.type == "cuda":
    scaler = torch.amp.GradScaler(device_type='cuda')  # Updated syntax
else:
    scaler = None  # Disable mixed precision on CPU

# Learning rate scheduler
scheduler = ReduceLROnPlateau(optimizer, mode='max', patience=2, factor=0.5)

"""### **Step 8: Training Loop**
We train the model for a fixed number of epochs. The training loop includes:
- Forward pass, loss computation, and backward pass.
- Mixed precision training if CUDA is available.
- Validation after each epoch to monitor performance.
- Saving the best model based on validation accuracy.

"""

# Training loop
best_val_acc = 0
num_epochs = 1

for epoch in range(num_epochs):
    # Training phase
    model.train()
    train_metrics = defaultdict(float)

    for batch_idx, (images, labels) in enumerate(train_loader):
        images, labels = images.to(device), labels.to(device)
        optimizer.zero_grad()

        if device.type == "cuda":
            with torch.autocast(device_type='cuda', dtype=torch.float16):
                logits = model(images)
                loss = criterion(logits, labels.unsqueeze(1))

            scaler.scale(loss).backward()
            scaler.step(optimizer)
            scaler.update()
        else:
            logits = model(images)
            loss = criterion(logits, labels.unsqueeze(1))
            loss.backward()
            optimizer.step()

        # Compute metrics
        preds = torch.sigmoid(logits)
        acc = ((preds > 0.5) == labels.unsqueeze(1)).float().mean()
        train_metrics['loss'] += loss.item()
        train_metrics['accuracy'] += acc.item()

        # Alert: Batch completed
        print(f"Epoch {epoch+1}/{num_epochs}, Batch {batch_idx+1}/{len(train_loader)} completed | Loss: {loss.item():.4f}, Acc: {acc.item():.2%}")

    # Validation phase
    model.eval()
    val_metrics = defaultdict(float)
    with torch.no_grad():
        for images, labels in val_loader:
            images, labels = images.to(device), labels.to(device)
            logits = model(images)
            loss = criterion(logits, labels.unsqueeze(1))
            preds = torch.sigmoid(logits)
            acc = ((preds > 0.5) == labels.unsqueeze(1)).float().mean()
            val_metrics['loss'] += loss.item()
            val_metrics['accuracy'] += acc.item()

    # Calculate averages
    train_loss = train_metrics['loss'] / len(train_loader)
    train_acc = train_metrics['accuracy'] / len(train_loader)
    val_loss = val_metrics['loss'] / len(val_loader)
    val_acc = val_metrics['accuracy'] / len(val_loader)

    # Scheduler step
    scheduler.step(val_acc)

    # Print epoch metrics
    print(f"\nEpoch {epoch+1}/{num_epochs} Summary:")
    print(f"Train Loss: {train_loss:.4f} | Train Acc: {train_acc:.2%}")
    print(f"Val Loss: {val_loss:.4f} | Val Acc: {val_acc:.2%}\n")

    # Save best model
    if val_acc > best_val_acc:
        torch.save(model.state_dict(), "best_model.pth")
        best_val_acc = val_acc
        print("New best model saved!\n")

"""### **Step 9: Visualization**
After training, we visualize some predictions to evaluate the model's performance qualitatively.

"""

# Load the best model
model.load_state_dict(torch.load("best_model.pth"))
model.eval()

# Visualize predictions
images, labels = next(iter(val_loader))
images, labels = images.to(device), labels.to(device)
with torch.no_grad():
    preds = torch.sigmoid(model(images))

plt.figure(figsize=(12, 6))
for i in range(8):
    plt.subplot(2, 4, i+1)
    plt.imshow(images[i].cpu().permute(1, 2, 0).clip(0, 1))
    plt.title(f"True: {labels[i].item()}\nPred: {preds[i].item():.2f}")
    plt.axis('off')
plt.tight_layout()
plt.show()

"""**⚠ Alert:** This training is for **test purposes only** and runs under **limited Google Colab resources** with a **restricted time frame**. Due to these constraints, training is set to **3 epochs** using **1% of the dataset**, with **small batch sizes** for efficiency. As a result, the model's performance will be **significantly lower** due to **fewer epochs and smaller batches**. Results may not reflect full-scale performance.
"""