48 lines
1.8 KiB
Python
48 lines
1.8 KiB
Python
|
|
import torch
|
|
import torch.nn as nn
|
|
from .Strategy import Strategy
|
|
from torch.utils.data import DataLoader
|
|
|
|
class LinearFiltration(Strategy):
|
|
def __init__(self, target_class_idx: int):
|
|
super().__init__()
|
|
self.target_class_idx = target_class_idx
|
|
|
|
def _run(self, model: nn.Module, forget_loader: DataLoader, retain_loader: DataLoader) -> 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)
|
|
# Filter the bias (if the layer uses one)
|
|
if model.fc.bias is not None:
|
|
b = model.fc.bias.data.clone()
|
|
# b is a 1D tensor of shape (num_classes),
|
|
# so we use torch.mv (matrix-vector multiplication) or unsqueeze it
|
|
sanitized_b = torch.mv(A, b)
|
|
model.fc.bias.copy_(sanitized_b)
|
|
|
|
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 = num_classes - 1
|
|
|
|
# The row of the forgotten class should average all other classes
|
|
for j in range(num_classes):
|
|
if j == forget_class:
|
|
# we zero the forget class
|
|
A[forget_class, j] = 0.0
|
|
else:
|
|
# and we distribute the output to the remaining
|
|
A[forget_class, j] = 1.0 / num_remaining
|
|
|
|
return A |