47 lines
1.8 KiB
Python
47 lines
1.8 KiB
Python
|
|
import torch.nn as nn
|
|
import time
|
|
import os
|
|
|
|
class Strategy:
|
|
"""Abstract base class for unlearning algorithms with automated, strategy-specific logging."""
|
|
|
|
def __init__(self):
|
|
# Dynamically set file name based on the class name (e.g., 'NormalizingLinearFiltration.txt')
|
|
self.strategy_name = self.__class__.__name__
|
|
self.log_file = f"reports/{self.strategy_name}_metrics.txt"
|
|
self._initialize_log_file()
|
|
|
|
def _initialize_log_file(self):
|
|
"""Creates a unique log file for this strategy with a header if it doesn't exist."""
|
|
if not os.path.exists(self.log_file):
|
|
with open(self.log_file, "w") as f:
|
|
f.write("execution_time_sec\n")
|
|
|
|
def log_metric(self, execution_time: float):
|
|
"""Appends the execution time to this strategy's specific file."""
|
|
with open(self.log_file, "a") as f:
|
|
f.write(f"{execution_time:.6f}\n")
|
|
|
|
def apply(self, model: nn.Module) -> nn.Module:
|
|
"""
|
|
Wraps the unlearning execution with automated timing and strategy-specific logging.
|
|
DO NOT override this method in subclasses. Override _run instead.
|
|
"""
|
|
start_time = time.perf_counter()
|
|
|
|
# Execute core unlearning logic
|
|
processed_model = self._run(model)
|
|
|
|
end_time = time.perf_counter()
|
|
execution_time = end_time - start_time
|
|
|
|
# Log to the strategy's specific file
|
|
self.log_metric(execution_time)
|
|
print(f"[{self.strategy_name}] Completed in {execution_time:.6f} seconds. Saved to {self.log_file}")
|
|
|
|
return processed_model
|
|
|
|
def _run(self, model: nn.Module) -> nn.Module:
|
|
"""Subclasses implement their core unlearning logic here."""
|
|
raise NotImplementedError |