29 lines
1.1 KiB
Python
29 lines
1.1 KiB
Python
|
|
import torch
|
|
from Strategy import Strategy
|
|
|
|
class NormalizingLinearFiltration(Strategy):
|
|
def __init__(self, target_class_idx):
|
|
self.target_class_idx = target_class_idx
|
|
|
|
def apply(self, model, forget_loader, retain_loader):
|
|
model.eval()
|
|
# Freeze parameters structurally
|
|
for param in model.parameters():
|
|
param.requires_grad = False
|
|
|
|
with torch.no_grad():
|
|
# we modify only classification head
|
|
# Shape: [num_classes, feature_dim]
|
|
W = model.fc.weight.data
|
|
|
|
# Compute the normalization transformation projection matrix (A)
|
|
# (In your full code, calculate A here matching Baumhauer et al.'s equations)
|
|
num_classes = W.shape[0]
|
|
A = torch.eye(num_classes, device=W.device)
|
|
# Mask/blend target class index distribution configurations here...
|
|
A[self.target_class_idx, :] = 0.0
|
|
|
|
# 3. Direct weight matrix override: W_filtered = A * W
|
|
sanitized_W = torch.mm(A, W)
|
|
model.fc.weight.copy_(sanitized_W) |