271 lines
9.6 KiB
Python
271 lines
9.6 KiB
Python
|
|
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)
|
|
return 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() |