| """Cofiber Threshold V3: spatial-reg variant with 3x3 depthwise convolution on regression features. |
| |
| Adds spatial context to box regression by letting each token see its 8 spatial |
| neighbors through a depthwise convolution. This directly addresses the |
| localization bottleneck in V2 (mAP@0.75 collapse) where each token predicts |
| box boundaries from its own features alone, with no access to neighboring tokens. |
| |
| Architecture delta from V2 (box32_92k): |
| - Add nn.Conv2d(32, 32, 3, padding=1, groups=32) after the regression hidden layer |
| - +320 params (32 channels * 9 kernel values + 32 bias) |
| - 91,960 total params (under NanoDet-m-0.5x head at 94K) |
| |
| Results vs V2: |
| - mAP@[0.5:0.95]: 5.9 -> 8.2 (+39%) |
| - mAP@0.50: 21.4 -> 26.1 (+22%) |
| - mAP@0.75: 1.3 -> 2.7 (+108%) |
| """ |
|
|
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
|
|
| NUM_CLASSES = 80 |
|
|
|
|
| def cofiber_decompose(f, n_scales): |
| cofibers = [] |
| residual = f |
| for _ in range(n_scales - 1): |
| omega = F.avg_pool2d(residual, 2) |
| sigma_omega = F.interpolate(omega, size=residual.shape[2:], mode="bilinear", align_corners=False) |
| cofibers.append(residual - sigma_omega) |
| residual = omega |
| cofibers.append(residual) |
| return cofibers |
|
|
|
|
| class CofiberThresholdV3(nn.Module): |
| """Cofiber decomposition + LayerNorm + prototype cls + spatial box regression. ~92K params.""" |
| name = "cofiber_threshold_spatialreg" |
| needs_intermediates = False |
|
|
| def __init__(self, feat_dim=768, num_classes=NUM_CLASSES, n_scales=3, reg_hidden=32): |
| super().__init__() |
| self.n_scales = n_scales |
| self.scale_norms = nn.ModuleList([nn.LayerNorm(feat_dim) for _ in range(n_scales)]) |
| self.prototypes = nn.Parameter(torch.randn(num_classes, feat_dim) * 0.01) |
| self.proto_bias = nn.Parameter(torch.zeros(num_classes)) |
| self.reg_hidden = nn.Linear(feat_dim, reg_hidden) |
| self.reg_act = nn.GELU() |
| |
| self.reg_spatial = nn.Conv2d(reg_hidden, reg_hidden, 3, padding=1, groups=reg_hidden) |
| self.reg_out = nn.Linear(reg_hidden, 4) |
| self.ctr_weight = nn.Parameter(torch.randn(1, feat_dim) * 0.01) |
| self.ctr_bias = nn.Parameter(torch.zeros(1)) |
| self.scale_params = nn.Parameter(torch.ones(n_scales)) |
|
|
| def forward(self, spatial, inter=None): |
| cofibers = cofiber_decompose(spatial, self.n_scales) |
| cls_l, reg_l, ctr_l = [], [], [] |
| for i, cof in enumerate(cofibers): |
| B, C, H, W = cof.shape |
| f = self.scale_norms[i](cof.permute(0, 2, 3, 1).reshape(-1, C)) |
| cls = (f @ self.prototypes.T + self.proto_bias).reshape(B, H, W, -1).permute(0, 3, 1, 2) |
| |
| reg_h = self.reg_act(self.reg_hidden(f)) |
| reg_h_spatial = reg_h.reshape(B, H, W, -1).permute(0, 3, 1, 2) |
| reg_h_spatial = self.reg_spatial(reg_h_spatial) |
| reg_h = reg_h_spatial.permute(0, 2, 3, 1).reshape(-1, reg_h.shape[-1]) |
| reg_raw = (self.reg_out(reg_h) * self.scale_params[i]).clamp(-10, 10) |
| reg = torch.exp(reg_raw).reshape(B, H, W, 4).permute(0, 3, 1, 2) |
| ctr = (f @ self.ctr_weight.T + self.ctr_bias).reshape(B, H, W, 1).permute(0, 3, 1, 2) |
| cls_l.append(cls) |
| reg_l.append(reg) |
| ctr_l.append(ctr) |
| return cls_l, reg_l, ctr_l |
|
|
