87 lines
3.4 KiB
Python
87 lines
3.4 KiB
Python
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_class_indices: torch.Tensor) -> torch.Tensor:
|
|
# 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 normally
|
|
x = self.layer4[0](x)
|
|
|
|
# -------------------------------------------------------------
|
|
# HERE IT IS: Save the structural skip connection (identity)
|
|
# BEFORE modifying features via block 1's convolutions
|
|
# -------------------------------------------------------------
|
|
identity = x
|
|
|
|
# Now enter layer4 block 1
|
|
x = self.layer4[1].conv1(x)
|
|
x = self.layer4[1].bn1(x)
|
|
x = self.layer4[1].relu(x)
|
|
|
|
# [Your Step 1 Masking Math happens right here...]
|
|
batch_alpha = self.alpha[target_class_indices]
|
|
mask = torch.sigmoid(batch_alpha).view(x.size(0), -1, 1, 1)
|
|
|
|
# Run the functional convolution
|
|
x = F.conv2d(
|
|
x,
|
|
weight=self.original_w,
|
|
bias=self.target_conv.bias,
|
|
stride=self.target_conv.stride,
|
|
padding=self.target_conv.padding
|
|
)
|
|
|
|
# Apply your WF-Net channel mask
|
|
x = x * mask
|
|
x = self.layer4[1].bn2(x)
|
|
|
|
# -------------------------------------------------------------
|
|
# HERE IT IS USED: Add the pristine identity back to the gated output
|
|
# -------------------------------------------------------------
|
|
x = self.layer4[1].relu(x + identity)
|
|
|
|
# Final Classification Head Sequence
|
|
x = self.avgpool(x)
|
|
x = torch.flatten(x, 1)
|
|
y_out = self.fc(x)
|
|
|
|
return y_out
|