diff --git a/Tune_new.py b/Tune_new.py index d8f155d..e9152f1 100644 --- a/Tune_new.py +++ b/Tune_new.py @@ -13,6 +13,8 @@ from unlearning.CertifiedRemoval import CertifiedRemoval from unlearning.CertifiedUnlearning import CertifiedUnlearning from unlearning.LinearFiltration import LinearFiltration from unlearning.WeightFiltration import WeightFiltration +from unlearning.WF import WeightF + # Global Hyperparameters CLASS_SIZE = 20 @@ -255,16 +257,16 @@ if __name__ == "__main__": target_class_index=i ) - weight_filtration = WeightFiltration( + weight_filtration = WeightF( #WeightFiltration( target_class_index=i, epochs=3, - lr=0.5, - gamma=150 + lr=0.05, + gamma=5 ) strategies = [ - certified_unlearning, - # weight_filtration, + # certified_unlearning, + weight_filtration, # linear_filtration ] diff --git a/unlearning/WeightFiltration.py b/unlearning/WeightFiltration.py index 3dd03eb..48a321e 100644 --- a/unlearning/WeightFiltration.py +++ b/unlearning/WeightFiltration.py @@ -59,20 +59,21 @@ class WeightFiltration(Strategy): temperature = 3.0 logits_f_scaled = outputs_f / temperature - loss_f = -torch.sum( - (torch.ones_like(logits_f_scaled) / num_classes) * torch.log_softmax(logits_f_scaled, dim=-1) - ) + # Compute uniform target entropy per-sample, then average over the batch + log_probs_f = torch.log_softmax(logits_f_scaled, dim=-1) + uniform_target = torch.ones_like(logits_f_scaled) / num_classes + loss_f = -torch.sum(uniform_target * log_probs_f, dim=-1).mean() + total_loss = loss_r + (self.gamma * loss_f) total_loss.backward() optimizer.step() - + t_loss_r += loss_r.item() t_loss_f += loss_f.item() steps += 1 - + print(f" Epoch {epoch+1}/{self.epochs} | Retain Loss: {t_loss_r/steps:.4f} | Forget Loss: {t_loss_f/steps:.4f}") - return wf_model @@ -110,15 +111,16 @@ class WeightFiltration(Strategy): ) # --- PERMANENT BAKING STEP --- - # Disconnect the dynamic parameter and freeze the optimal gated state permanently into the architecture - with torch.no_grad(): + # Grab the alpha mask vector for the forgotten class and cast to 4D tensor shape final_mask = torch.sigmoid(wf_model.alpha[self.target_class_index]).view(-1, 1, 1, 1) - target_conv.weight.copy_(original_weights * final_mask) - # Re-enable model parameters for downstream evaluation processing + # Apply filter masking permanently back onto the base layer + target_conv.weight.copy_(original_weights * final_mask) + + # Unfreeze architecture parameters for evaluations downstream for p in model.parameters(): p.requires_grad = True - - print(f">> Permanently altered {out_channels} convolutional filters in layer4 via WF-Net.") - return model \ No newline at end of file + + print(f">> Permanently altered {out_channels} convolutional filters in layer4 via WF-Net.") + return model diff --git a/unlearning/wf/WF_Net.py b/unlearning/wf/WF_Net.py new file mode 100644 index 0000000..a5dd690 --- /dev/null +++ b/unlearning/wf/WF_Net.py @@ -0,0 +1,86 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F + +class WF_Net(nn.Module): + """ + Implements Poppi et al.'s WF Model structure. + Wraps a pre-trained ResNet-18 and dynamically applies + weight-space gating matrix multiplication during the forward step. + """ + def __init__(self, original_model: nn.Module, num_classes: int): + super().__init__() + # Extract the sequence of blocks/layers L from the original model + self.conv1 = original_model.conv1 + self.bn1 = original_model.bn1 + self.relu = original_model.relu + self.maxpool = original_model.maxpool + + self.layer1 = original_model.layer1 + self.layer2 = original_model.layer2 + self.layer3 = original_model.layer3 + self.layer4 = original_model.layer4 + + self.avgpool = original_model.avgpool + self.fc = original_model.fc + + # Target layer for filtering: layer4 block 1 conv2 + # We extract its static tensor data out of the autograd parameter pool + self.target_conv = self.layer4[1].conv2 + self.original_w = nn.Parameter(self.target_conv.weight.data.clone().detach(), requires_grad=False) + + # Require: Alpha gating matrix. Shape: (num_classes, out_channels) + # Initialized to 1.5 as per Poppi et al.'s verbatim specification + out_channels = self.original_w.shape[0] + #self.alpha = nn.Parameter(torch.ones(num_classes, out_channels) * 1.5) + self.alpha = nn.Parameter(torch.ones(num_classes, out_channels)) + + def forward(self, x: torch.Tensor, target_unlearn_class: int) -> torch.Tensor: + """ + Implements Algorithm 1: General forward step of a WF model + Inputs: + x: Input tensor (Xin) + target_unlearn_class: The class index we are actively filtering out (Yunl) + """ + # 1. Run through early sequence of layers undisturbed + x = self.maxpool(self.relu(self.bn1(self.conv1(x)))) + x = self.layer1(x) + x = self.layer2(x) + x = self.layer3(x) + + # Run layer4 block 0 and block 1 conv1 normally + x = self.layer4[0](x) + + identity = x + + + + x = self.layer4[1].conv1(x) + x = self.layer4[1].bn1(x) + x = self.layer4[1].relu(x) + + # 2. CORE WF-NET MATH: w_hat_l <- alpha_l[Yunl] ⊙ w_l + # Extract 1D vector for target class and reshape to (out_channels, 1, 1, 1) for 4D convolution broadcasting + mask = torch.sigmoid(self.alpha[target_unlearn_class]).view(-1, 1, 1, 1) + w_hat = self.original_w * mask + + # 3. Pass gated weights straight to functional forward pass: l(Xi, w_hat_l) + x = F.conv2d( + x, + weight=w_hat, + bias=self.target_conv.bias, + stride=self.target_conv.stride, + padding=self.target_conv.padding + ) + x = self.layer4[1].bn2(x) + + # Handle residual shortcut skip connection manually since we opened up block 1 + # In ResNet-18 layer4, block 1 has no downsample shortcut layer; it's a direct identity add + x = self.layer4[1].relu(x + identity) + + # 4. Final Classification Head Sequence + x = self.avgpool(x) + x = torch.flatten(x, 1) + y_out = self.fc(x) + + return y_out