127 lines
4.9 KiB
Python
127 lines
4.9 KiB
Python
import torch
|
||
import torch.nn as nn
|
||
import torch.nn.functional as F
|
||
import torch.optim as optim
|
||
from torch.utils.data import DataLoader
|
||
import numpy as np
|
||
from sklearn.metrics import classification_report
|
||
from architectures.Model import Model
|
||
|
||
class WF_Module(nn.Module):
|
||
"""
|
||
Pure PyTorch Neural Network module graph.
|
||
Keeps parameter registration and autograd tracking separate from
|
||
the framework's high-level Model abstractions to prevent recursion collisions.
|
||
"""
|
||
def __init__(self, original_model: nn.Module, num_classes: int):
|
||
super().__init__()
|
||
|
||
self.original_model = original_model
|
||
|
||
# Target layer for weight filtering (layer4 block 1 conv2 or conv3 depending on arch)
|
||
last_layer = original_model.layer4[1]
|
||
|
||
# Some versions are limited to 2 convolutional layers
|
||
if hasattr(last_layer, "conv3"):
|
||
self.target_conv = last_layer.conv3
|
||
else:
|
||
self.target_conv = last_layer.conv2
|
||
|
||
# Completely freeze the original ResNet parameters
|
||
for param in self.parameters():
|
||
param.requires_grad = False
|
||
|
||
# Initialize the alpha parameter matrix (Rows = Classes, Cols = Channels)
|
||
out_channels = self.target_conv.weight.shape[0]
|
||
self.alpha = nn.Parameter(torch.full((num_classes, out_channels), 3.0))
|
||
|
||
'''
|
||
Poppi et_al's Single-shot multiclass unlearning.
|
||
This calculation happens only once to generate the mask. once the mask is generated,
|
||
Unlearning and remembering becomes a matter of switching gates on and off.
|
||
'''
|
||
def forward(self, x: torch.Tensor, target_class_indices: torch.Tensor) -> torch.Tensor:
|
||
# we linearly loop through layers 1 to 4[block 1] (for ResNet)
|
||
# for i in M_{|L|} do l <- l[i]
|
||
x = self.original_model.maxpool(self.original_model.relu(self.original_model.bn1(self.original_model.conv1(x))))
|
||
x = self.original_model.layer1(x)
|
||
x = self.original_model.layer2(x)
|
||
x = self.original_model.layer3(x)
|
||
x = self.original_model.layer4[0](x)
|
||
|
||
# The second block execute its internal transformations natively
|
||
# This handles conv1->conv2 (ResNet18) or conv1->conv2->conv3 (ResNet50) automatically!
|
||
# Xi+1 <- l(Xi, ˆwl)
|
||
x = self.original_model.layer4[1](x)
|
||
|
||
# Apply mask dynamically to the completed block feature map
|
||
# wl <- αl[Yunl] ⊙ ˆwl
|
||
batch_alpha = self.alpha[target_class_indices]
|
||
mask = torch.sigmoid(batch_alpha).view(x.size(0), -1, 1, 1)
|
||
x = x * mask
|
||
|
||
# Remaining standard head steps
|
||
x = self.original_model.avgpool(x)
|
||
x = torch.flatten(x, 1)
|
||
# so here we are returning the output logits
|
||
# the result of classification is then
|
||
# argmax(x)
|
||
return self.original_model.fc(x)
|
||
|
||
|
||
|
||
class WF_Net_Model(Model):
|
||
def __init__(self, device, size, original_model: nn.Module, target_class_index: int):
|
||
self.device = device
|
||
self.size = size
|
||
self.wf_module = WF_Module(
|
||
original_model = original_model,
|
||
num_classes = size
|
||
).to(self.device)
|
||
|
||
# this index indicates which row of the mask should be active (gate closed).
|
||
self.target_class_index = target_class_index
|
||
self.model = self.wf_module
|
||
|
||
def get(self):
|
||
return self.wf_module
|
||
|
||
'''
|
||
We override the evaluate method from the base class,
|
||
because how we evaluate is different here from that of a normal torch nn.Module object
|
||
|
||
'''
|
||
def evaluate(self, loader, mode="eval"):
|
||
|
||
self.wf_module.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)
|
||
|
||
# we apply the filter
|
||
gate_signals = torch.full((inputs.size(0),), self.target_class_index, dtype=torch.long, device=self.device)
|
||
|
||
# pass prediction through the filter
|
||
outputs = self.wf_module(inputs, target_class_indices=gate_signals)
|
||
|
||
# return argmax(x)
|
||
_, 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(all_labels)
|
||
|
||
print(f"Test Accuracy: {accuracy:.2f}%")
|
||
print(classification_report(all_labels, all_preds, labels=classes, zero_division=0))
|
||
report = classification_report(all_labels, all_preds, labels=classes, output_dict=True, zero_division=0)
|
||
|
||
return accuracy, report
|
||
|
||
def eval(self):
|
||
"""Safely intercept any fallback base class calls targeting .eval()"""
|
||
self.wf_module.eval()
|