Finetuning/architectures/Model.py

from abc import ABC, abstractmethod
import torch
import torch.nn as nn
import torch.optim as optim
import time
import numpy as np
from sklearn.metrics import classification_report
from pathlib import Path
#from unlearning.Strategy import Strategy
import copy
from torch.optim.lr_scheduler import CosineAnnealingLR
import Util

class Model(ABC):
    # need to add a weight decay here
    def __init__(self, device, size):
        self.device = device
        self.size = size
        self.model = self.get().to(self.device)

    @abstractmethod
    def get(self):
        pass

    '''
   
    '''
    def train(self, epochs, loader, rate, mode = "retrain"):
        criterion = nn.CrossEntropyLoss()
        optimizer = optim.Adam(filter(lambda p: p.requires_grad, self.model.parameters()), lr=rate, weight_decay=0.1)
        
        scheduler = CosineAnnealingLR(optimizer, T_max=epochs, eta_min=1e-6)
        # to save reports
        # file_path = Path(f"{mode}/{self.__class__.__name__.lower()}/time_metrics.txt")
        # Util._initialize_log_file(file_path)
        

        print(f"Starting training on {self.device}...")
        # start_time = time.time()
        # training phase
        self.model.train()

        for epoch in range(epochs):
            total_loss = 0.0
            for inputs, labels in loader:
                inputs, labels = inputs.to(self.device), labels.to(self.device)
                # zero param gradients
                optimizer.zero_grad()
                # forward pass
                outputs = self.model(inputs)
                # comoutew loss
                loss = criterion(outputs, labels)
                # backward pass
                loss.backward()
                optimizer.step()
                total_loss += loss.item()
            scheduler.step()
        
            print(f"Epoch {epoch+1}/{epochs} | Loss: {total_loss / len(loader):.4f}")
        #end_time = time.time()
        #execution_time = end_time - start_time
        #Util.log_metric(log_file=file_path, execution_time=execution_time)
        if self.device.type == 'cuda': torch.cuda.synchronize()
        print(f"Training complete.")


    def save(self, filename=None):
        
        save_dir = Path("trained_models")
        save_dir.mkdir(parents=True, exist_ok=True)

        # Filename (Default to class name if not provided)
        if filename is None:
            filename = f"{self.__class__.__name__.lower()}.pth"
        
        if not filename.endswith('.pth'):
            filename += '.pth'

        save_path = save_dir / filename
        torch.save(self.model.state_dict(), save_path)
        print(f'Model saved to {save_path}')



    def load(self, arch):
 
        file_path = Path("trained_models") / f'{arch.name.lower()}.pth'
        # does file exist
        if not file_path.exists():
            raise FileNotFoundError(f'No checkpoint found at: {file_path}')
            
        # Load the weights
        state_dict = torch.load(file_path, map_location=self.device, weights_only=True)
        
        self.model.load_state_dict(state_dict)
        self.model.to(self.device)
        print(f'Model loaded from {file_path}')

 

    def evaluate(self, loader, mode="eval"):
        """
        Evaluates the model, prints terminal reports, and routes metrics to 
        a file logger based on the current context mode.
        """
        
        self.model.eval()
        all_preds, all_labels = [], []
        print(f"\nEvaluating Domain: [{mode}]...")
        
        with torch.no_grad():
            for inputs, labels in loader:
                inputs, labels = inputs.to(self.device), labels.to(self.device)
                outputs = self.model(inputs)
                _, predicted = torch.max(outputs, 1)
                all_preds.extend(predicted.cpu().numpy())
                all_labels.extend(labels.cpu().numpy())

        # Extract only the active classes evaluated in this loader slice
        classes = sorted(list(set(all_labels)))
        accuracy = 100 * (np.array(all_preds) == np.array(all_labels)).sum() / len(all_labels)
        
        print(f"Test Accuracy: {accuracy:.2f}%")
        
        # 1. Print standard text report to terminal
        print(classification_report(all_labels, all_preds, labels=classes, zero_division=0))
        
        # 2. Extract structured dictionary metrics
        report_dict = classification_report(
            all_labels, 
            all_preds, 
            labels=classes, 
            output_dict=True, 
            zero_division=0
        )
        
        # 3. Delegate file tracking to isolated helper method
        #self._log_to_csv(mode, accuracy,report_dict)
        return accuracy, report_dict



    # factory
    @staticmethod
    def create(arch, device, size):
        print(f'>> MODEL ARCHITECTURE  >> {arch.name}.')

        match arch:

            # ResNet18
            case Architecture.RESNET18: 
                from architectures.ResNet18 import ResNet18
                return ResNet18(device, size)
            
            # ResNet34
            case Architecture.RESNET34: 
                from architectures.ResNet34 import ResNet34
                return ResNet34(device, size)
            
            # ResNet50
            case Architecture.RESNET50:
                from architectures.ResNet50 import ResNet50
                return ResNet50(device, size)
            
            # INCEPTION
            case Architecture.INCEPTION:
                from architectures.Inception import Inception
                return Inception(device, size)
            
            # DENSENET121
            case Architecture.DENSENET121:
                from architectures.DenseNet121 import DenseNet121
                return DenseNet121(device, size)
            # googleNet
            case Architecture.GOOGLENET:
                from architectures.GoogleNet import GoogleNet
                return GoogleNet(device, size)
            # EfficientNet
            case Architecture.EFFICIENTNET:
                from architectures.EfficentNet import EfficientNet
                return EfficientNet(device, size)
            #ShuffleNet
            case Architecture.SHUFFLENET:
                from architectures.ShuffleNet import ShuffleNet
                return ShuffleNet(device, size)
            # wide ResNet
            case Architecture.WIDE_RESNET:
                from architectures.WideResNet import WideResNet
                return WideResNet(device, size)
            case _:
                raise ValueError(f"Unknown model: {arch}")

# model architectures
from enum import Enum, auto

class Architecture(Enum):
    RESNET18 = auto()
    RESNET50 = auto()
    RESNET34 = auto()
    INCEPTION = auto()
    DENSENET121 = auto()
    GOOGLENET = auto()
    EFFICIENTNET = auto()
    SHUFFLENET = auto()
    WIDE_RESNET = auto()