169 lines
6.7 KiB
Python
169 lines
6.7 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):
|
|
def __init__(self, original_model: nn.Module, num_classes: int, arch_enum):
|
|
super().__init__()
|
|
# If your model classes contain the raw inner torch model under an attribute,
|
|
# extract it. Otherwise, use it directly.
|
|
self.original_model = getattr(original_model, "model", original_model)
|
|
|
|
# Freeze the original model parameters completely
|
|
for param in self.original_model.parameters():
|
|
param.requires_grad = False
|
|
|
|
# Target layer discovery using your clean Enum contract
|
|
self.target_layer = self._deduce_target_layer(self.original_model, arch_enum)
|
|
|
|
# Derive channel dimensions dynamically from the deduced layer
|
|
out_channels = self._extract_channels(self.target_layer, self.original_model)
|
|
|
|
# Initialize alpha parameter matrix (Rows = Classes, Cols = Channels)
|
|
self.alpha = nn.Parameter(torch.full((num_classes, out_channels), 3.0))
|
|
self._current_target_indices = None
|
|
|
|
def _deduce_target_layer(self, model: nn.Module, arch_enum) -> nn.Module:
|
|
"""
|
|
Scans the architecture topology to target the final deep feature extraction block
|
|
right before global pooling/classification using strict Enum configurations.
|
|
"""
|
|
match arch_enum:
|
|
# --- RESNET FAMILY ---
|
|
case arch_enum.RESNET18 | arch_enum.RESNET34 | arch_enum.RESNET50 | arch_enum.WIDE_RESNET:
|
|
return model.layer4[-1]
|
|
|
|
# --- GOOGLENET ---
|
|
case arch_enum.GOOGLENET:
|
|
return model.inception5b
|
|
|
|
# --- INCEPTION V3 ---
|
|
case arch_enum.INCEPTION:
|
|
return model.Mixed_7c
|
|
|
|
# --- DENSENET 121 ---
|
|
case arch_enum.DENSENET121:
|
|
return model.features.norm5
|
|
|
|
# --- EFFICIENTNET ---
|
|
case arch_enum.EFFICIENTNET:
|
|
return model.features[-1]
|
|
|
|
# --- SHUFFLENET ---
|
|
case arch_enum.SHUFFLENET:
|
|
return model.conv5
|
|
|
|
case _:
|
|
# Robust Fallback Strategy
|
|
target = None
|
|
for module in model.modules():
|
|
if isinstance(module, nn.Conv2d):
|
|
target = module
|
|
if target is not None:
|
|
return target
|
|
raise RuntimeError(f"Could not locate filtration anchor for Enum target: {arch_enum}")
|
|
|
|
def _extract_channels(self, target_layer: nn.Module, model: nn.Module) -> int:
|
|
"""Helper to determine channel depth across varied layers types."""
|
|
if hasattr(target_layer, "out_channels"):
|
|
return target_layer.out_channels
|
|
if hasattr(target_layer, "num_features"):
|
|
return target_layer.num_features
|
|
if hasattr(target_layer, "weight"):
|
|
return target_layer.weight.shape[0]
|
|
|
|
# Classifier fallback mapping
|
|
if hasattr(model, "fc"):
|
|
return model.fc.in_features
|
|
if hasattr(model, "classifier"):
|
|
if isinstance(model.classifier, nn.Linear):
|
|
return model.classifier.in_features
|
|
if isinstance(model.classifier, nn.Sequential):
|
|
return model.classifier[0].in_features
|
|
return 512
|
|
|
|
def _filtration_hook(self, module: nn.Module, hook_input: tuple, hook_output: torch.Tensor) -> torch.Tensor:
|
|
if self._current_target_indices is None:
|
|
return hook_output
|
|
|
|
batch_alpha = self.alpha[self._current_target_indices]
|
|
|
|
if len(hook_output.shape) == 4:
|
|
mask = torch.sigmoid(batch_alpha).view(hook_output.size(0), -1, 1, 1)
|
|
else:
|
|
mask = torch.sigmoid(batch_alpha).view(hook_output.size(0), -1)
|
|
|
|
return hook_output * mask
|
|
|
|
def forward(self, x: torch.Tensor, target_class_indices: torch.Tensor) -> torch.Tensor:
|
|
self._current_target_indices = target_class_indices
|
|
hook_handle = self.target_layer.register_forward_hook(self._filtration_hook)
|
|
try:
|
|
logits = self.original_model(x)
|
|
finally:
|
|
hook_handle.remove()
|
|
self._current_target_indices = None
|
|
return logits
|
|
|
|
|
|
class WF_Net_Model(Model):
|
|
def __init__(self, device, size, original_model: nn.Module, target_class_index: int, arch):
|
|
self.device = device
|
|
self.size = size
|
|
self.wf_module = WF_Module(
|
|
arch_enum=arch,
|
|
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()
|