| """ |
| ALPINE: Adaptive Localization for Parameter- and Sample-Efficient Few-Shot Learning |
| Official Model Architecture & Checkpoint Loader |
| |
| This module defines: |
| - Canonical EXP-F3 (22,249 parameters): ALPINE_CIFAR, ALPINE_Native |
| - Optional Variant EXP-F3-35k (34,917 parameters): ALPINE_35k_CIFAR, ALPINE_35k_Native |
| - Gabor edge-energy guided windowed patch locator |
| - Convenient `load_alpine_model` checkpoint loader |
| """ |
|
|
| import math |
| import os |
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
|
|
| try: |
| from .irfe_p1_canonical import ( |
| GaborPreprocess, |
| PatchEncoderCIFARP1, |
| SubPatchEncoderCIFARP1, |
| BoundaryEncoderCIFARP1, |
| PatchEncoderNativeP1, |
| SubPatchEncoderNativeP1, |
| BoundaryEncoderNativeP1 |
| ) |
| except ImportError: |
| from irfe_p1_canonical import ( |
| GaborPreprocess, |
| PatchEncoderCIFARP1, |
| SubPatchEncoderCIFARP1, |
| BoundaryEncoderCIFARP1, |
| PatchEncoderNativeP1, |
| SubPatchEncoderNativeP1, |
| BoundaryEncoderNativeP1 |
| ) |
|
|
|
|
| def get_exp_p1_base_centers(): |
| """Returns canonical base 5-patch center coordinates in normalized [-1, 1] range.""" |
| return torch.tensor([ |
| [-0.5161, -0.5161], |
| [-0.5161, 0.5161], |
| [ 0.0000, 0.0000], |
| [ 0.5161, -0.5161], |
| [ 0.5161, 0.5161] |
| ], dtype=torch.float32) |
|
|
|
|
| class WideWindowAdaptivePatchLocator(nn.Module): |
| """ |
| Windowed Adaptive Patch Locator (EXP-F3). |
| Dynamically displaces patch sampling centers toward salient features |
| within a Gaussian spatial window around canonical base centers, |
| guided by Gabor edge-energy response E(y, x). |
| """ |
| def __init__(self, base_centers, window_frac=0.50, temperature=1.0): |
| super().__init__() |
| self.register_buffer("base_centers", base_centers) |
| self.num_patches = base_centers.shape[0] |
| self.window_frac = window_frac |
| self.temperature = nn.Parameter(torch.tensor(temperature, dtype=torch.float32)) |
|
|
| def forward(self, energy_map): |
| B, _, H, W = energy_map.shape |
| ys = torch.linspace(-1, 1, H, device=energy_map.device) |
| xs = torch.linspace(-1, 1, W, device=energy_map.device) |
| gy, gx = torch.meshgrid(ys, xs, indexing='ij') |
|
|
| flat_energy = energy_map.squeeze(1) |
|
|
| centers = [] |
| for k in range(self.num_patches): |
| cy0 = self.base_centers[k, 0] |
| cx0 = self.base_centers[k, 1] |
| |
| dist_sq = (gy - cy0)**2 + (gx - cx0)**2 |
| log_window = -dist_sq / (2 * (self.window_frac**2) + 1e-6) |
| |
| weighted_energy = flat_energy / (self.temperature.abs() + 1e-4) + log_window |
| w = F.softmax(weighted_energy.view(B, -1), dim=-1) |
| cy = (w * gy.reshape(-1)).sum(-1) |
| cx = (w * gx.reshape(-1)).sum(-1) |
| centers.append(torch.stack([cx, cy], dim=-1)) |
| return torch.stack(centers, dim=1) |
|
|
|
|
| def extract_patches_grid_sample(x, centers, patch_scale=0.5): |
| """Bilinearly extracts image patches around adaptive centers using grid_sample.""" |
| B, C, H, W = x.shape |
| patch_h = int(round(H * patch_scale)) |
| patch_w = int(round(W * patch_scale)) |
|
|
| patches = [] |
| for k in range(centers.size(1)): |
| cx = centers[:, k, 0] |
| cy = centers[:, k, 1] |
| |
| theta = torch.zeros(B, 2, 3, device=x.device, dtype=x.dtype) |
| theta[:, 0, 0] = patch_scale |
| theta[:, 1, 1] = patch_scale |
| theta[:, 0, 2] = cx |
| theta[:, 1, 2] = cy |
|
|
| grid = F.affine_grid(theta, torch.Size([B, C, patch_h, patch_w]), align_corners=False) |
| p = F.grid_sample(x, grid, align_corners=False, mode='bilinear', padding_mode='reflection') |
| patches.append(p) |
| return patches |
|
|
|
|
| |
| |
| |
|
|
| class ALPINE_CIFAR(nn.Module): |
| """ |
| Canonical ALPINE / EXP-F3 architecture for CIFAR-FS (32x32). |
| Total Trainable Parameters: 22,249. |
| """ |
| def __init__(self, window_frac=0.50, embed_dim=16, num_heads=4): |
| super().__init__() |
| self.embed_dim = embed_dim |
| base_centers = get_exp_p1_base_centers() |
| self.locator = WideWindowAdaptivePatchLocator(base_centers, window_frac=window_frac) |
| self.encoder = PatchEncoderCIFARP1(out_dim=embed_dim) |
| self.sub_encoder = SubPatchEncoderCIFARP1(out_dim=embed_dim) |
| self.boundary_encoder = BoundaryEncoderCIFARP1(out_dim=embed_dim) |
| self.sub_fusion = nn.Sequential( |
| nn.Linear(embed_dim * 4, embed_dim), nn.LayerNorm(embed_dim), nn.ReLU() |
| ) |
| self.rel_proj = nn.Sequential( |
| nn.Linear(embed_dim * 2, embed_dim), nn.LayerNorm(embed_dim), nn.ReLU() |
| ) |
| self.mha = nn.MultiheadAttention(embed_dim, num_heads, batch_first=True) |
| self.ln_attn = nn.LayerNorm(embed_dim) |
| self.ws_query = nn.Parameter(torch.randn(1, 1, embed_dim) * 0.02) |
| self.gate_net = nn.Sequential( |
| nn.Linear(embed_dim * 2, 16), nn.ReLU(), nn.Linear(16, embed_dim * 2) |
| ) |
|
|
| def _split_boundaries(self, x): |
| return (x[:,:, 8:16, 8:16], x[:,:, 8:16, 16:24], |
| x[:,:, 16:24, 8:16], x[:,:, 16:24, 16:24]) |
|
|
| def _encode_patch(self, p): |
| v_c = self.encoder(p) |
| sp = [p[:,:, r:r+8, c:c+8] for r in [0, 8] for c in [0, 8]] |
| v_f = self.sub_fusion(torch.cat([self.sub_encoder(s) for s in sp], dim=1)) |
| return v_c + v_f |
|
|
| def _rel(self, va, vb): |
| return self.rel_proj(torch.cat([va - vb, va * vb], dim=1)) |
|
|
| def extract_with_rel_tokens(self, x): |
| B = x.size(0) |
| gabor_out = self.encoder.gabor(x) |
| energy_map = gabor_out[:, 3:7].abs().sum(dim=1, keepdim=True) |
| centers = self.locator(energy_map) |
|
|
| raw_patches = extract_patches_grid_sample(x, centers, patch_scale=0.5) |
| patches = [self._encode_patch(p) for p in raw_patches] |
| edges = [self._rel(patches[i], patches[j]) for i in range(5) for j in range(i+1, 5)] |
| boundaries = [self.boundary_encoder(b) for b in self._split_boundaries(x)] |
| |
| tokens = torch.stack(patches + edges + boundaries, dim=1) |
| attn, _ = self.mha(tokens, tokens, tokens) |
| refined = self.ln_attn(tokens + attn) |
| ws = self.ws_query.expand(B, -1, -1) |
| scores = torch.bmm(ws, refined.transpose(1, 2)) / math.sqrt(self.embed_dim) |
| W = torch.bmm(F.softmax(scores, dim=-1), refined).squeeze(1) |
| |
| patch_pool = refined[:, :5, :].mean(dim=1) |
| g = torch.cat([patch_pool, W], dim=1) |
| gate = 2.0 * torch.sigmoid(self.gate_net(g)) - 1.0 |
| out_feat= g + g * gate |
|
|
| rel_tokens = torch.stack(edges, dim=1) |
| return out_feat, centers, rel_tokens |
|
|
| def extract(self, x): |
| """Extracts fixed-dimensional feature representations for few-shot metric classification.""" |
| feat, _, _ = self.extract_with_rel_tokens(x) |
| return feat |
|
|
| def compute_prototypes(self, sx, sy, n=5): |
| """Computes support prototypes for n-way classification.""" |
| f = self.extract(sx) |
| return torch.stack([f[sy == c].mean(0) for c in range(n)]) |
|
|
| def predict_proto(self, qx, protos): |
| """Computes negative squared Euclidean distance between query representations and prototypes.""" |
| return -(torch.cdist(self.extract(qx), protos) ** 2) |
|
|
|
|
| class ALPINE_Native(nn.Module): |
| """ |
| Canonical ALPINE / EXP-F3 architecture for MiniImageNet Native (84x84). |
| Total Trainable Parameters: 22,249. |
| """ |
| def __init__(self, window_frac=0.50, embed_dim=16, num_heads=4): |
| super().__init__() |
| self.embed_dim = embed_dim |
| base_centers = get_exp_p1_base_centers() |
| self.locator = WideWindowAdaptivePatchLocator(base_centers, window_frac=window_frac) |
| self.encoder = PatchEncoderNativeP1(out_dim=embed_dim) |
| self.sub_encoder = SubPatchEncoderNativeP1(out_dim=embed_dim) |
| self.boundary_encoder = BoundaryEncoderNativeP1(out_dim=embed_dim) |
| self.sub_fusion = nn.Sequential( |
| nn.Linear(embed_dim * 4, embed_dim), nn.LayerNorm(embed_dim), nn.ReLU() |
| ) |
| self.rel_proj = nn.Sequential( |
| nn.Linear(embed_dim * 2, embed_dim), nn.LayerNorm(embed_dim), nn.ReLU() |
| ) |
| self.mha = nn.MultiheadAttention(embed_dim, num_heads, batch_first=True) |
| self.ln_attn = nn.LayerNorm(embed_dim) |
| self.ws_query = nn.Parameter(torch.randn(1, 1, embed_dim) * 0.02) |
| self.gate_net = nn.Sequential( |
| nn.Linear(embed_dim * 2, 16), nn.ReLU(), nn.Linear(16, embed_dim * 2) |
| ) |
|
|
| def _split_boundaries(self, x): |
| return (x[:,:, 18:48, 18:48], x[:,:, 18:48, 36:66], |
| x[:,:, 36:66, 18:48], x[:,:, 36:66, 36:66]) |
|
|
| def _encode_patch(self, p): |
| v_c = self.encoder(p) |
| sp = [p[:,:, r:r+24, c:c+24] for r in [0, 24] for c in [0, 24]] |
| v_f = self.sub_fusion(torch.cat([self.sub_encoder(s) for s in sp], dim=1)) |
| return v_c + v_f |
|
|
| def _rel(self, va, vb): |
| return self.rel_proj(torch.cat([va - vb, va * vb], dim=1)) |
|
|
| def extract_with_rel_tokens(self, x): |
| B = x.size(0) |
| gabor_out = self.encoder.gabor(x) |
| energy_map = gabor_out[:, 3:7].abs().sum(dim=1, keepdim=True) |
| centers = self.locator(energy_map) |
|
|
| raw_patches = extract_patches_grid_sample(x, centers, patch_scale=48.0/84.0) |
| patches = [self._encode_patch(p) for p in raw_patches] |
| edges = [self._rel(patches[i], patches[j]) for i in range(5) for j in range(i+1, 5)] |
| boundaries = [self.boundary_encoder(b) for b in self._split_boundaries(x)] |
| |
| tokens = torch.stack(patches + edges + boundaries, dim=1) |
| attn, _ = self.mha(tokens, tokens, tokens) |
| refined = self.ln_attn(tokens + attn) |
| ws = self.ws_query.expand(B, -1, -1) |
| scores = torch.bmm(ws, refined.transpose(1, 2)) / math.sqrt(self.embed_dim) |
| W = torch.bmm(F.softmax(scores, dim=-1), refined).squeeze(1) |
| |
| patch_pool = refined[:, :5, :].mean(dim=1) |
| g = torch.cat([patch_pool, W], dim=1) |
| gate = 2.0 * torch.sigmoid(self.gate_net(g)) - 1.0 |
| out_feat= g + g * gate |
|
|
| rel_tokens = torch.stack(edges, dim=1) |
| return out_feat, centers, rel_tokens |
|
|
| def extract(self, x): |
| feat, _, _ = self.extract_with_rel_tokens(x) |
| return feat |
|
|
| def compute_prototypes(self, sx, sy, n=5): |
| f = self.extract(sx) |
| return torch.stack([f[sy == c].mean(0) for c in range(n)]) |
|
|
| def predict_proto(self, qx, protos): |
| return -(torch.cdist(self.extract(qx), protos) ** 2) |
|
|
|
|
| |
| IRFEExpF3_CIFAR = ALPINE_CIFAR |
| IRFEExpF3_Native = ALPINE_Native |
|
|
|
|
| |
| |
| |
|
|
| class PatchEncoderCIFARGeneric(nn.Module): |
| def __init__(self, c1, c2, out_dim): |
| super().__init__() |
| self.gabor = GaborPreprocess(ksize=5) |
| self.conv1 = nn.Conv2d(7, c1, 3, padding=1) |
| self.conv2 = nn.Conv2d(c1, c2, 3, padding=1) |
| self.pool = nn.MaxPool2d(2, 2) |
| self.fc = nn.Linear(c2 * 4 * 4, out_dim) |
| self.ln = nn.LayerNorm(out_dim) |
|
|
| def forward(self, x): |
| x = self.gabor(x) |
| x = self.pool(F.relu(self.conv1(x))) |
| x = self.pool(F.relu(self.conv2(x))) |
| return self.ln(self.fc(x.view(x.size(0), -1))) |
|
|
|
|
| class SubPatchEncoderCIFARGeneric(nn.Module): |
| def __init__(self, sub_c, out_dim): |
| super().__init__() |
| self.gabor = GaborPreprocess(ksize=5) |
| self.conv = nn.Conv2d(7, sub_c, 3, padding=1) |
| self.pool = nn.AdaptiveAvgPool2d((4, 4)) |
| self.fc = nn.Linear(sub_c * 4 * 4, out_dim) |
|
|
| def forward(self, x): |
| x = self.gabor(x) |
| return F.relu(self.fc(self.pool(F.relu(self.conv(x))).view(x.size(0), -1))) |
|
|
|
|
| class BoundaryEncoderCIFARGeneric(nn.Module): |
| def __init__(self, sub_c, out_dim): |
| super().__init__() |
| self.gabor = GaborPreprocess(ksize=5) |
| self.conv = nn.Conv2d(7, sub_c, 3, padding=1) |
| self.pool = nn.AdaptiveAvgPool2d((4, 4)) |
| self.fc = nn.Linear(sub_c * 4 * 4, out_dim) |
|
|
| def forward(self, x): |
| x = self.gabor(x) |
| return F.relu(self.fc(self.pool(F.relu(self.conv(x))).view(x.size(0), -1))) |
|
|
|
|
| class PatchEncoderNativeGeneric(nn.Module): |
| def __init__(self, c1, c2, out_dim): |
| super().__init__() |
| self.gabor = GaborPreprocess(ksize=5) |
| self.conv1 = nn.Conv2d(7, c1, 3, padding=1) |
| self.conv2 = nn.Conv2d(c1, c2, 3, padding=1) |
| self.pool = nn.MaxPool2d(2, 2) |
| self.adap = nn.AdaptiveAvgPool2d((4, 4)) |
| self.fc = nn.Linear(c2 * 4 * 4, out_dim) |
| self.ln = nn.LayerNorm(out_dim) |
|
|
| def forward(self, x): |
| x = self.gabor(x) |
| x = self.pool(F.relu(self.conv1(x))) |
| x = self.pool(F.relu(self.conv2(x))) |
| x = self.adap(x) |
| return self.ln(self.fc(x.view(x.size(0), -1))) |
|
|
|
|
| class SubPatchEncoderNativeGeneric(nn.Module): |
| def __init__(self, sub_c, out_dim): |
| super().__init__() |
| self.gabor = GaborPreprocess(ksize=5) |
| self.conv = nn.Conv2d(7, sub_c, 3, padding=1) |
| self.pool = nn.AdaptiveAvgPool2d((4, 4)) |
| self.fc = nn.Linear(sub_c * 4 * 4, out_dim) |
|
|
| def forward(self, x): |
| x = self.gabor(x) |
| return F.relu(self.fc(self.pool(F.relu(self.conv(x))).view(x.size(0), -1))) |
|
|
|
|
| class BoundaryEncoderNativeGeneric(nn.Module): |
| def __init__(self, sub_c, out_dim): |
| super().__init__() |
| self.gabor = GaborPreprocess(ksize=5) |
| self.conv = nn.Conv2d(7, sub_c, 3, padding=1) |
| self.pool = nn.AdaptiveAvgPool2d((4, 4)) |
| self.fc = nn.Linear(sub_c * 4 * 4, out_dim) |
|
|
| def forward(self, x): |
| x = self.gabor(x) |
| return F.relu(self.fc(self.pool(F.relu(self.conv(x))).view(x.size(0), -1))) |
|
|
|
|
| class ALPINE_35k_CIFAR(nn.Module): |
| """ |
| Optional Variant: ALPINE / EXP-F3-35k for CIFAR-FS (32x32). |
| Total Trainable Parameters: 34,917. |
| """ |
| def __init__(self, c1=15, c2=19, sub_c=8, gate_h=12, embed_dim=32, window_frac=0.50, num_heads=4): |
| super().__init__() |
| self.embed_dim = embed_dim |
| base_centers = get_exp_p1_base_centers() |
| self.locator = WideWindowAdaptivePatchLocator(base_centers, window_frac=window_frac) |
| self.encoder = PatchEncoderCIFARGeneric(c1, c2, out_dim=embed_dim) |
| self.sub_encoder = SubPatchEncoderCIFARGeneric(sub_c, out_dim=embed_dim) |
| self.boundary_encoder = BoundaryEncoderCIFARGeneric(sub_c, out_dim=embed_dim) |
| self.sub_fusion = nn.Sequential( |
| nn.Linear(embed_dim * 4, embed_dim), nn.LayerNorm(embed_dim), nn.ReLU() |
| ) |
| self.rel_proj = nn.Sequential( |
| nn.Linear(embed_dim * 2, embed_dim), nn.LayerNorm(embed_dim), nn.ReLU() |
| ) |
| self.mha = nn.MultiheadAttention(embed_dim, num_heads, batch_first=True) |
| self.ln_attn = nn.LayerNorm(embed_dim) |
| self.ws_query = nn.Parameter(torch.randn(1, 1, embed_dim) * 0.02) |
| self.gate_net = nn.Sequential( |
| nn.Linear(embed_dim * 2, gate_h), nn.ReLU(), nn.Linear(gate_h, embed_dim * 2) |
| ) |
|
|
| def _split_boundaries(self, x): |
| return (x[:,:, 8:16, 8:16], x[:,:, 8:16, 16:24], |
| x[:,:, 16:24, 8:16], x[:,:, 16:24, 16:24]) |
|
|
| def _encode_patch(self, p): |
| v_c = self.encoder(p) |
| sp = [p[:,:, r:r+8, c:c+8] for r in [0, 8] for c in [0, 8]] |
| v_f = self.sub_fusion(torch.cat([self.sub_encoder(s) for s in sp], dim=1)) |
| return v_c + v_f |
|
|
| def _rel(self, va, vb): |
| return self.rel_proj(torch.cat([va - vb, va * vb], dim=1)) |
|
|
| def extract_with_rel_tokens(self, x): |
| B = x.size(0) |
| gabor_out = self.encoder.gabor(x) |
| energy_map = gabor_out[:, 3:7].abs().sum(dim=1, keepdim=True) |
| centers = self.locator(energy_map) |
|
|
| raw_patches = extract_patches_grid_sample(x, centers, patch_scale=0.5) |
| patches = [self._encode_patch(p) for p in raw_patches] |
| edges = [self._rel(patches[i], patches[j]) for i in range(5) for j in range(i+1, 5)] |
| boundaries = [self.boundary_encoder(b) for b in self._split_boundaries(x)] |
| |
| tokens = torch.stack(patches + edges + boundaries, dim=1) |
| attn, _ = self.mha(tokens, tokens, tokens) |
| refined = self.ln_attn(tokens + attn) |
| ws = self.ws_query.expand(B, -1, -1) |
| scores = torch.bmm(ws, refined.transpose(1, 2)) / math.sqrt(self.embed_dim) |
| W = torch.bmm(F.softmax(scores, dim=-1), refined).squeeze(1) |
| |
| patch_pool = refined[:, :5, :].mean(dim=1) |
| g = torch.cat([patch_pool, W], dim=1) |
| gate = 2.0 * torch.sigmoid(self.gate_net(g)) - 1.0 |
| out_feat= g + g * gate |
| rel_tokens = torch.stack(edges, dim=1) |
| return out_feat, centers, rel_tokens |
|
|
| def extract(self, x): |
| feat, _, _ = self.extract_with_rel_tokens(x) |
| return feat |
|
|
| def compute_prototypes(self, sx, sy, n=5): |
| f = self.extract(sx) |
| return torch.stack([f[sy == c].mean(0) for c in range(n)]) |
|
|
| def predict_proto(self, qx, protos): |
| return -(torch.cdist(self.extract(qx), protos) ** 2) |
|
|
|
|
| class ALPINE_35k_Native(nn.Module): |
| """ |
| Optional Variant: ALPINE / EXP-F3-35k for MiniImageNet Native (84x84). |
| Total Trainable Parameters: 34,917. |
| """ |
| def __init__(self, c1=15, c2=19, sub_c=8, gate_h=12, embed_dim=32, window_frac=0.50, num_heads=4): |
| super().__init__() |
| self.embed_dim = embed_dim |
| base_centers = get_exp_p1_base_centers() |
| self.locator = WideWindowAdaptivePatchLocator(base_centers, window_frac=window_frac) |
| self.encoder = PatchEncoderNativeGeneric(c1, c2, out_dim=embed_dim) |
| self.sub_encoder = SubPatchEncoderNativeGeneric(sub_c, out_dim=embed_dim) |
| self.boundary_encoder = BoundaryEncoderNativeGeneric(sub_c, out_dim=embed_dim) |
| self.sub_fusion = nn.Sequential( |
| nn.Linear(embed_dim * 4, embed_dim), nn.LayerNorm(embed_dim), nn.ReLU() |
| ) |
| self.rel_proj = nn.Sequential( |
| nn.Linear(embed_dim * 2, embed_dim), nn.LayerNorm(embed_dim), nn.ReLU() |
| ) |
| self.mha = nn.MultiheadAttention(embed_dim, num_heads, batch_first=True) |
| self.ln_attn = nn.LayerNorm(embed_dim) |
| self.ws_query = nn.Parameter(torch.randn(1, 1, embed_dim) * 0.02) |
| self.gate_net = nn.Sequential( |
| nn.Linear(embed_dim * 2, gate_h), nn.ReLU(), nn.Linear(gate_h, embed_dim * 2) |
| ) |
|
|
| def _split_boundaries(self, x): |
| return (x[:,:, 18:48, 18:48], x[:,:, 18:48, 36:66], |
| x[:,:, 36:66, 18:48], x[:,:, 36:66, 36:66]) |
|
|
| def _encode_patch(self, p): |
| v_c = self.encoder(p) |
| sp = [p[:,:, r:r+24, c:c+24] for r in [0, 24] for c in [0, 24]] |
| v_f = self.sub_fusion(torch.cat([self.sub_encoder(s) for s in sp], dim=1)) |
| return v_c + v_f |
|
|
| def _rel(self, va, vb): |
| return self.rel_proj(torch.cat([va - vb, va * vb], dim=1)) |
|
|
| def extract_with_rel_tokens(self, x): |
| B = x.size(0) |
| gabor_out = self.encoder.gabor(x) |
| energy_map = gabor_out[:, 3:7].abs().sum(dim=1, keepdim=True) |
| centers = self.locator(energy_map) |
|
|
| raw_patches = extract_patches_grid_sample(x, centers, patch_scale=48.0/84.0) |
| patches = [self._encode_patch(p) for p in raw_patches] |
| edges = [self._rel(patches[i], patches[j]) for i in range(5) for j in range(i+1, 5)] |
| boundaries = [self.boundary_encoder(b) for b in self._split_boundaries(x)] |
| |
| tokens = torch.stack(patches + edges + boundaries, dim=1) |
| attn, _ = self.mha(tokens, tokens, tokens) |
| refined = self.ln_attn(tokens + attn) |
| ws = self.ws_query.expand(B, -1, -1) |
| scores = torch.bmm(ws, refined.transpose(1, 2)) / math.sqrt(self.embed_dim) |
| W = torch.bmm(F.softmax(scores, dim=-1), refined).squeeze(1) |
| |
| patch_pool = refined[:, :5, :].mean(dim=1) |
| g = torch.cat([patch_pool, W], dim=1) |
| gate = 2.0 * torch.sigmoid(self.gate_net(g)) - 1.0 |
| out_feat= g + g * gate |
| rel_tokens = torch.stack(edges, dim=1) |
| return out_feat, centers, rel_tokens |
|
|
| def extract(self, x): |
| feat, _, _ = self.extract_with_rel_tokens(x) |
| return feat |
|
|
| def compute_prototypes(self, sx, sy, n=5): |
| f = self.extract(sx) |
| return torch.stack([f[sy == c].mean(0) for c in range(n)]) |
|
|
| def predict_proto(self, qx, protos): |
| return -(torch.cdist(self.extract(qx), protos) ** 2) |
|
|
|
|
| |
| IRFEExpF3_35k_CIFAR = ALPINE_35k_CIFAR |
| IRFEExpF3_35k_Native = ALPINE_35k_Native |
|
|
|
|
| def load_alpine_model(checkpoint_path, model_type="canonical", dataset="cifar", device="cpu"): |
| """ |
| Convenience loader for ALPINE model checkpoints. |
| |
| Args: |
| checkpoint_path (str): Path to .pt checkpoint file. |
| model_type (str): 'canonical' (22,249 params) or 'variant-35k' (34,917 params). |
| dataset (str): 'cifar' (32x32 images) or 'mini' (84x84 images). |
| device (str or torch.device): Device to load model onto. |
| |
| Returns: |
| nn.Module: Loaded ALPINE model ready in eval mode. |
| """ |
| device = torch.device(device) |
| model_type = model_type.lower() |
| dataset = dataset.lower() |
|
|
| if "35k" in model_type: |
| if dataset in ["cifar", "cifar-fs", "cifar_fs"]: |
| model = ALPINE_35k_CIFAR() |
| else: |
| model = ALPINE_35k_Native() |
| else: |
| if dataset in ["cifar", "cifar-fs", "cifar_fs"]: |
| model = ALPINE_CIFAR() |
| else: |
| model = ALPINE_Native() |
|
|
| checkpoint = torch.load(checkpoint_path, map_location=device, weights_only=False) |
| state_dict = checkpoint["model_state_dict"] if isinstance(checkpoint, dict) and "model_state_dict" in checkpoint else checkpoint |
| model.load_state_dict(state_dict) |
| model.to(device) |
| model.eval() |
| return model |
|
|
