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 class Model(ABC): 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): criterion = nn.CrossEntropyLoss() optimizer = optim.Adam(filter(lambda p: p.requires_grad, self.model.parameters()), lr=rate) print(f"Starting training on {self.device}...") start_time = time.time() 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) optimizer.zero_grad() outputs = self.model(inputs) loss = criterion(outputs, labels) loss.backward() optimizer.step() total_loss += loss.item() print(f"Epoch {epoch+1}/{epochs} | Loss: {total_loss / len(loader):.4f}") if self.device.type == 'cuda': torch.cuda.synchronize() print(f"Training completed in: {time.time() - start_time:.2f}s") def evaluate(self, loader): self.model.eval() all_preds, all_labels = [], [] print("\nEvaluating...") 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()) classes = sorted(list(set(all_labels))) accuracy = 100 * (np.array(all_preds) == np.array(all_labels)).sum() / len(classes) print(f"Test Accuracy: {accuracy:.2f}%") print(classification_report(all_labels, all_preds, labels=classes, zero_division=0)) def save(self, filename=None): save_dir = Path("trained_models") save_dir.mkdir(parents=True, exist_ok=True) # Determine 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 unlearn(self, strategy: Strategy, forget_loader, retain_loader): """ Executes a targeted unlearning strategy and profiles efficiency """ print(f"Executing: {strategy.__class__.__name__}...") start_time = time.time() # Delegate the actual algorithmic weight/logit manipulation to the strategy strategy.apply(self.model, forget_loader, retain_loader) elapsed_time = time.time() - start_time print(f"{strategy.__class__.__name__} completed in {elapsed_time:.4f} seconds.") return elapsed_time 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) def _log_to_csv(self, mode, accuracy, report_dict): """Handles directory structures, file setups, and distinct CSV column formatting.""" arch_name = self.__class__.__name__.lower() save_dir = Path("reports") save_dir.mkdir(parents=True, exist_ok=True) csv_path = save_dir / f"{arch_name}-{mode}.csv" file_exists = csv_path.exists() ''' # Structure payload and headers based on evaluation slice type if mode == "forget": headers = ["accuracy", "precision", "recall", "f1-score"] target_cls_str = str(classes[0]) metrics = report_dict[target_cls_str] row = [ f"{accuracy / 100.0:.4f}", f"{metrics['precision']:.4f}", f"{metrics['recall']:.4f}", f"{metrics['f1-score']:.4f}" ] else: headers = [ "accuracy", "macro_precision", "macro_recall", "macro_f1", "weighted_precision", "weighted_recall", "weighted_f1" ] row = [ f"{accuracy / 100.0:.4f}", f"{report_dict['macro avg']['precision']:.4f}", f"{report_dict['macro avg']['recall']:.4f}", f"{report_dict['macro avg']['f1-score']:.4f}", f"{report_dict['weighted avg']['precision']:.4f}", f"{report_dict['weighted avg']['recall']:.4f}", f"{report_dict['weighted avg']['f1-score']:.4f}" ]''' headers = [ "accuracy", "macro_precision", "macro_recall", "macro_f1", "weighted_precision", "weighted_recall", "weighted_f1" ] row = [ f"{accuracy / 100.0:.4f}", f"{report_dict['macro avg']['precision']:.4f}", f"{report_dict['macro avg']['recall']:.4f}", f"{report_dict['macro avg']['f1-score']:.4f}", f"{report_dict['weighted avg']['precision']:.4f}", f"{report_dict['weighted avg']['recall']:.4f}", f"{report_dict['weighted avg']['f1-score']:.4f}" ] with open(csv_path, "a") as f: if not file_exists: f.write(",".join(headers) + "\n") f.write(",".join(row) + "\n") print(f">> Direct CSV metrics appended to {csv_path}") # Using the factory patern here @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) # 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() INCEPTION = auto() DENSENET121 = auto() GOOGLENET = auto() EFFICIENTNET = auto() SHUFFLENET = auto() WIDE_RESNET = auto()