module.py

import math
import torch
import torchvision.transforms as T
from os import path
from torch.utils.data import DataLoader, WeightedRandomSampler
from torch.optim import AdamW
from torch.optim.lr_scheduler import CosineAnnealingLR
from torch.nn import CrossEntropyLoss
from torchmetrics.functional import accuracy
from timm import create_model, list_models
from timm.models.vision_transformer import VisionTransformer
from torchvision.datasets import ImageFolder

from utils import AverageMeter
from lightning import LightningDataModule, LightningModule
from huggingface_hub import PyTorchModelHubMixin, login
import torch.nn as nn
from lora import LoRA_qkv


PRE_SIZE = (256, 256)
IMG_SIZE = (224, 224)

STATS = dict(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
DATASET_DIRECTORY = path.join(path.dirname(__file__), "datasets")
CHECKPOINT_DIRECTORY = path.join(path.dirname(__file__), "checkpoints")

TRANSFORMS = {
    "train": T.Compose([
        T.Resize(PRE_SIZE),
        T.RandomCrop(IMG_SIZE),
        T.ToTensor(),
        T.Normalize(**STATS)
    ]),
    "val": T.Compose([
        T.Resize(PRE_SIZE),
        T.CenterCrop(IMG_SIZE),
        T.ToTensor(),
        T.Normalize(**STATS)
    ])
}



class myDataModule(LightningDataModule):
    """
    Lightning DataModule for loading and preparing the image dataset.

    Args:
        ds_name (str): Name of the dataset directory.
        batch_size (int): Batch size for data loaders.
        num_workers (int): Number of workers for data loaders.
    """
    def __init__(self, ds_name: str = "deities", batch_size: int = 32, num_workers: int = 8):
        super(myDataModule, self).__init__()

        self.ds_path = path.join(DATASET_DIRECTORY, ds_name)
        assert path.exists(self.ds_path), f"Dataset {ds_name} not found in {DATASET_DIRECTORY}."

        self.ds_name = ds_name
        self.batch_size = batch_size
        self.num_workers = num_workers


    def setup(self, stage=None):
        if stage == "fit" or stage is None:
            self.train_ds = ImageFolder(root=path.join(self.ds_path, 'train'), transform=TRANSFORMS['train'])
            self.val_ds = ImageFolder(root=path.join(self.ds_path, 'val'), transform=TRANSFORMS['val'])
            # Number of classes
            self.num_classes = len(self.train_ds.classes)          
    

    def train_dataloader(self) -> DataLoader:
        # Weighted Random sampler for imbalanced dataset
        class_samples = [0] * self.num_classes
        for _, (_, label) in enumerate(self.train_ds):
            class_samples[label] += 1
        weights = [1.0 / class_samples[label] for _, label in self.train_ds]
        self.sampler = WeightedRandomSampler(weights, len(weights), replacement=True)
        return DataLoader(dataset=self.train_ds, batch_size=self.batch_size, 
                          sampler=self.sampler, num_workers=self.num_workers, persistent_workers=True)


    def val_dataloader(self) -> DataLoader:
        return DataLoader(dataset=self.val_ds, batch_size=self.batch_size, 
                          shuffle=False, num_workers=self.num_workers, persistent_workers=True)
    



class myModule(LightningModule, PyTorchModelHubMixin):
    """
    Lightning Module for training and evaluating the Image classification model.

    Args:
        model_name (str): Name of the Vision Transformer model.
        num_classes (int): Number of classes in the dataset.
        freeze_flag (bool): Flag to freeze the base model parameters.
        use_lora (bool): Flag to use LoRA (Local Rank Adaptation) for fine-tuning.
        rank (int): Rank for LoRA if use_lora is True.
        learning_rate (float): Learning rate for the optimizer.
        weight_decay (float): Weight decay for the optimizer.
        push_to_hf (bool): Flag to push model to Huggingface Hub.
        commit_message (str): Commit message
        repo_id (str): Huggingface repo id
    """
    def __init__(self, 
                 model_name: str = "vit_tiny_patch16_224", 
                 num_classes: int = 25,
                 freeze_flag: bool = True,
                 use_lora: bool = False, 
                 rank: int = None, 
                 learning_rate: float = 3e-4, 
                 weight_decay: float = 2e-5,
                 push_to_hf: bool = True,
                 commit_message: str = "my model",
                 repo_id: str = "Yegiiii/ideityfy"
        ):
    
        super(myModule, self).__init__()
        self.save_hyperparameters()
        self.model_name = model_name
        self.num_classes = num_classes
        self.freeze_flag = freeze_flag
        self.rank = rank
        self.use_lora = use_lora
        self.learning_rate = learning_rate
        self.weight_decay = weight_decay
        self.push_to_hf = push_to_hf
        self.commit_message = commit_message
        self.repo_id = repo_id
        
        assert model_name in list_models(), f"Timm model name {model_name} not available."
        timm_model = create_model(model_name, pretrained=True)
        assert isinstance(timm_model, VisionTransformer), f"{model_name} not a Vision Transformer."
        self.model = timm_model

        if freeze_flag:
            # Freeze the Timm model parameters
            self.freeze()

        if use_lora:
            # Add LoRA matrices to the Timm model
            assert freeze_flag, "Set freeze_flag to True for using LoRA fine-tuning."
            assert rank, "Rank can't be None."
            # self.model = LoRA_VisionTransformer(self.model, rank)
            self.add_lora()

        self.model.reset_classifier(num_classes)

        # Loss function
        self.criterion = CrossEntropyLoss()

        # Validation metrics
        self.top1_acc = AverageMeter()
        self.top3_acc = AverageMeter()
        self.top5_acc = AverageMeter()


    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.model(x)
    
    
    def on_fit_start(self) -> None:
        num_classes = self.trainer.datamodule.num_classes
        assert num_classes == self.num_classes, \
        f"Number of classes provided in the argument ({self.num_classes}) is not matching \
         the number of classes in the dataset ({num_classes})."


    def on_fit_end(self) -> None:
        if self.push_to_hf:
            login()
            self.push_to_hub(repo_id=self.repo_id, commit_message=self.commit_message)


    def configure_optimizers(self):
        optimizer = AdamW(params=filter(lambda param: param.requires_grad, self.model.parameters()), 
                          lr=self.learning_rate, weight_decay=self.weight_decay)
        
        scheduler = CosineAnnealingLR(optimizer, self.trainer.max_epochs, 1e-6)
        return ([optimizer], [scheduler])


    def shared_step(self, x: torch.Tensor, y: torch.Tensor):
        logits = self(x)
        loss = self.criterion(logits, y)   
        return logits, loss 


    def training_step(self, batch, batch_idx) -> torch.Tensor:
        x, y = batch
        _, loss = self.shared_step(x, y)

        self.log("train_loss", loss, prog_bar=True, logger=True, on_epoch=True)
        return loss


    def validation_step(self, batch, batch_idx) -> dict:
        x, y = batch
        logits, loss = self.shared_step(x, y)

        self.top1_acc(
            val=accuracy(logits, y, average="weighted", top_k=1, num_classes=self.num_classes))
        self.top3_acc(
            val=accuracy(logits, y, average="weighted", top_k=3, num_classes=self.num_classes))
        self.top5_acc(
            val=accuracy(logits, y, average="weighted", top_k=5, num_classes=self.num_classes))

        metric_dict = {
            "val_loss": loss, 
            "top1_acc": self.top1_acc.avg, 
            "top3_acc": self.top3_acc.avg, 
            "top5_acc": self.top5_acc.avg
        }
        
        self.log_dict(metric_dict, prog_bar=True, logger=True, on_epoch=True)
        return  metric_dict

    
    def on_validation_epoch_end(self) -> None:
        self.top1_acc.reset()
        self.top3_acc.reset()
        self.top5_acc.reset()


    def add_lora(self):
        self.w_As = []
        self.w_Bs = []

        for _, blk in enumerate(self.model.blocks):
            w_qkv_linear = blk.attn.qkv
            self.dim = w_qkv_linear.in_features
            lora_a_linear_q = nn.Linear(self.dim, self.rank, bias=False)
            lora_b_linear_q = nn.Linear(self.rank, self.dim, bias=False)
            lora_a_linear_v = nn.Linear(self.dim, self.rank, bias=False)
            lora_b_linear_v = nn.Linear(self.rank, self.dim, bias=False)
            self.w_As.append(lora_a_linear_q)
            self.w_Bs.append(lora_b_linear_q)
            self.w_As.append(lora_a_linear_v)
            self.w_Bs.append(lora_b_linear_v)
            blk.attn.qkv = LoRA_qkv(w_qkv_linear, lora_a_linear_q, 
                                    lora_b_linear_q, lora_a_linear_v, lora_b_linear_v)

        for w_A in self.w_As:
            nn.init.kaiming_uniform_(w_A.weight, a=math.sqrt(5))
        for w_B in self.w_Bs:
            nn.init.zeros_(w_B.weight)



if __name__ == "__main__":
    # from torchinfo import summary
    
    # module = myModule(freeze_flag=False)
    # summary(module, (1, 3, 224, 224))

    from datasets import load_dataset

    dataset = load_dataset("Yegiiii/deities")
    print(dataset)