48 lines
1.7 KiB
Python
48 lines
1.7 KiB
Python
|
|
import torch
|
|
import torch.nn as nn
|
|
from .Strategy import Strategy
|
|
|
|
class LinearFiltration(Strategy):
|
|
def __init__(self, target_class_idx: int):
|
|
super().__init__() # Automatically configures 'NormalizingLinearFiltration_metrics.txt'
|
|
self.target_class_idx = target_class_idx
|
|
|
|
def _run(self, model: nn.Module) -> nn.Module:
|
|
model.eval()
|
|
for param in model.parameters():
|
|
param.requires_grad = False
|
|
|
|
with torch.no_grad():
|
|
W = model.fc.weight.data.clone()
|
|
num_classes = W.shape[0]
|
|
|
|
A = self._calculate_filtration_matrix(num_classes, self.target_class_idx, W.device)
|
|
sanitized_W = torch.mm(A, W)
|
|
model.fc.weight.copy_(sanitized_W)
|
|
|
|
return model
|
|
|
|
'''@staticmethod
|
|
def _calculate_filtration_matrix(num_classes: int, forget_class: int, device: torch.device) -> torch.Tensor:
|
|
A = torch.eye(num_classes, device=device)
|
|
num_remaining_classes = num_classes - 1
|
|
for j in range(num_classes):
|
|
if j == forget_class:
|
|
A[forget_class, j] = 0.0
|
|
else:
|
|
A[forget_class, j] = 1.0 / num_remaining_classes
|
|
return A'''
|
|
@staticmethod
|
|
def _calculate_filtration_matrix(num_classes: int, forget_class: int, device: torch.device) -> torch.Tensor:
|
|
A = torch.eye(num_classes, device=device)
|
|
num_remaining = num_classes - 1
|
|
|
|
# The row of the forgotten class should average the inputs of all other classes
|
|
for j in range(num_classes):
|
|
if j == forget_class:
|
|
A[forget_class, j] = 0.0
|
|
else:
|
|
A[forget_class, j] = 1.0 / num_remaining
|
|
|
|
return A |