Sync Coast-style multi-scale fusion (part 2)

marine_dual_dinov3_backbone.py

@@ -55,6 +55,163 @@ def _set_module_trainable(module: nn.Module, trainable: bool) -> None:
         param.requires_grad = trainable
 
 
+class CrossScaleAttention(nn.Module):
+    """Use ViT tokens as queries to align one ConvNeXt feature scale."""
+
+    def __init__(self, dim: int, num_heads: int = 8, mlp_ratio: int = 4) -> None:
+        super().__init__()
+        self.q_norm = nn.LayerNorm(dim)
+        self.kv_norm = nn.LayerNorm(dim)
+        self.attn = nn.MultiheadAttention(embed_dim=dim, num_heads=num_heads, batch_first=True)
+        self.norm = nn.LayerNorm(dim)
+        self.ffn = nn.Sequential(
+            nn.Linear(dim, dim * mlp_ratio),
+            nn.GELU(),
+            nn.Linear(dim * mlp_ratio, dim),
+        )
+
+    def forward(self, vit_tokens: Tensor, conv_tokens: Tensor) -> Tensor:
+        aligned, _ = self.attn(
+            query=self.q_norm(vit_tokens),
+            key=self.kv_norm(conv_tokens),
+            value=self.kv_norm(conv_tokens),
+            need_weights=False,
+        )
+        x = vit_tokens + aligned
+        return x + self.ffn(self.norm(x))
+
+
+class CoastStyleMultiScaleFusion(nn.Module):
+    """CoastGPT-inspired heterogeneous multi-scale fusion for dense heads."""
+
+    def __init__(
+        self,
+        vit_dim: int,
+        convnext_dims: tuple[int, int, int, int],
+        out_channels: int = 256,
+        semantic_grid_size: tuple[int, int] | None = None,
+        detail_grid_size: tuple[int, int] = (28, 28),
+        num_heads: int = 8,
+    ) -> None:
+        super().__init__()
+        self.vit_dim = int(vit_dim)
+        self.out_channels = int(out_channels)
+        self.semantic_grid_size = semantic_grid_size
+        self.detail_grid_size = detail_grid_size
+
+        c4_dim, c8_dim, c16_dim, c32_dim = [int(v) for v in convnext_dims]
+        self.c4_proj = nn.Sequential(nn.Linear(c4_dim, vit_dim), nn.LayerNorm(vit_dim))
+        self.c16_proj = nn.Sequential(nn.Linear(c16_dim, vit_dim), nn.LayerNorm(vit_dim))
+        self.c32_proj = nn.Sequential(nn.Linear(c32_dim, vit_dim), nn.LayerNorm(vit_dim))
+        self.c4_attn = CrossScaleAttention(vit_dim, num_heads=num_heads)
+        self.c16_attn = CrossScaleAttention(vit_dim, num_heads=num_heads)
+        self.c32_attn = CrossScaleAttention(vit_dim, num_heads=num_heads)
+        self.semantic_logits = nn.Parameter(torch.zeros(4))
+        self.semantic_compress = nn.Sequential(
+            nn.Linear(vit_dim * 4, vit_dim),
+            nn.LayerNorm(vit_dim),
+            nn.GELU(),
+            nn.Linear(vit_dim, out_channels),
+            nn.LayerNorm(out_channels),
+        )
+        self.detail_proj = nn.Sequential(
+            nn.Linear(c8_dim, vit_dim),
+            nn.LayerNorm(vit_dim),
+            nn.GELU(),
+            nn.Linear(vit_dim, out_channels),
+            nn.LayerNorm(out_channels),
+        )
+        self.detail_to_semantic = nn.Sequential(
+            nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1, groups=out_channels, bias=False),
+            nn.Conv2d(out_channels, out_channels, kernel_size=1, bias=False),
+            nn.GroupNorm(1, out_channels),
+            nn.GELU(),
+        )
+        self.fusion_gate = nn.Sequential(
+            nn.Conv2d(out_channels * 2, out_channels, kernel_size=1),
+            nn.GELU(),
+            nn.Conv2d(out_channels, out_channels, kernel_size=1),
+            nn.Sigmoid(),
+        )
+        self.output_smooth = nn.Sequential(
+            nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1, bias=False),
+            nn.GroupNorm(1, out_channels),
+            nn.GELU(),
+        )
+        self._last_fusion_weights: Tensor | None = None
+
+    @staticmethod
+    def _add_2d_sincos_pos_embed(tokens: Tensor, h: int, w: int) -> Tensor:
+        device, dtype = tokens.device, tokens.dtype
+        channels = tokens.shape[-1]
+        quarter = max(channels // 4, 1)
+        yy, xx = torch.meshgrid(
+            torch.arange(h, device=device),
+            torch.arange(w, device=device),
+            indexing="ij",
+        )
+        omega = torch.arange(quarter, device=device, dtype=dtype)
+        omega = 1.0 / (10000 ** (omega / quarter))
+        yy = yy.reshape(-1, 1).to(dtype) * omega
+        xx = xx.reshape(-1, 1).to(dtype) * omega
+        pos = torch.cat([torch.sin(yy), torch.cos(yy), torch.sin(xx), torch.cos(xx)], dim=1)
+        if pos.shape[1] < channels:
+            pad = torch.zeros(pos.shape[0], channels - pos.shape[1], device=device, dtype=dtype)
+            pos = torch.cat([pos, pad], dim=1)
+        return tokens + pos[:, :channels].unsqueeze(0)
+
+    @classmethod
+    def _feature_to_tokens(cls, feat: Tensor, target_size: tuple[int, int] | None = None) -> Tensor:
+        if target_size is not None and feat.shape[-2:] != target_size:
+            feat = F.interpolate(feat, size=target_size, mode="bilinear", align_corners=False)
+        _, _, h, w = feat.shape
+        tokens = feat.flatten(2).transpose(1, 2).contiguous()
+        return cls._add_2d_sincos_pos_embed(tokens, h, w)
+
+    @staticmethod
+    def _tokens_to_map(tokens: Tensor, h: int, w: int) -> Tensor:
+        return tokens.transpose(1, 2).reshape(tokens.shape[0], tokens.shape[-1], h, w).contiguous()
+
+    def forward(self, vit_map: Tensor, convnext_feats: tuple[Tensor, Tensor, Tensor, Tensor]) -> dict[str, Tensor]:
+        c4, c8, c16, c32 = convnext_feats
+        if self.semantic_grid_size is not None and vit_map.shape[-2:] != self.semantic_grid_size:
+            vit_map = F.interpolate(vit_map, size=self.semantic_grid_size, mode="bilinear", align_corners=False)
+        h, w = vit_map.shape[-2:]
+        vit_tokens = self._feature_to_tokens(vit_map)
+
+        f4 = self.c4_attn(vit_tokens, self.c4_proj(self._feature_to_tokens(c4)))
+        f16 = self.c16_attn(vit_tokens, self.c16_proj(self._feature_to_tokens(c16)))
+        f32 = self.c32_attn(vit_tokens, self.c32_proj(self._feature_to_tokens(c32)))
+
+        weights = torch.softmax(self.semantic_logits, dim=0)
+        semantic_tokens = torch.cat(
+            [
+                vit_tokens * weights[0],
+                f4 * weights[1],
+                f16 * weights[2],
+                f32 * weights[3],
+            ],
+            dim=-1,
+        )
+        semantic_tokens = self.semantic_compress(semantic_tokens)
+        semantic = self._tokens_to_map(semantic_tokens, h, w)
+
+        detail_tokens = self.detail_proj(self._feature_to_tokens(c8, target_size=self.detail_grid_size))
+        detail = self._tokens_to_map(detail_tokens, self.detail_grid_size[0], self.detail_grid_size[1])
+        detail_down = F.interpolate(detail, size=semantic.shape[-2:], mode="bilinear", align_corners=False)
+        detail_down = self.detail_to_semantic(detail_down)
+
+        gate = self.fusion_gate(torch.cat([semantic, detail_down], dim=1))
+        fused = self.output_smooth(semantic + gate * detail_down)
+        self._last_fusion_weights = weights.detach()
+        return {
+            "fused": fused,
+            "semantic": semantic,
+            "detail": detail,
+            "fusion_weights": weights,
+        }
+
+
 class DINOv3ViTConvNeXtBackbone(nn.Module):
     """Dual-branch DINOv3 backbone with explicit partial unfreezing.
 
@@ -88,30 +245,21 @@ class DINOv3ViTConvNeXtBackbone(nn.Module):
         )
         self.vit_channels = int(getattr(self.vit, "embed_dim"))
         self.convnext_channels = int(getattr(self.convnext, "embed_dim"))
+        self.convnext_dims = tuple(int(v) for v in getattr(self.convnext, "embed_dims"))
         self.out_channels = out_channels
 
-        self.vit_proj = nn.Sequential(
-            nn.Conv2d(self.vit_channels, out_channels, kernel_size=1, bias=False),
-            nn.BatchNorm2d(out_channels),
-            nn.GELU(),
-        )
-        self.convnext_proj = nn.Sequential(
-            nn.Conv2d(self.convnext_channels, out_channels, kernel_size=1, bias=False),
-            nn.BatchNorm2d(out_channels),
-            nn.GELU(),
-        )
-        self.fusion = nn.Sequential(
-            nn.Conv2d(out_channels * 2, out_channels, kernel_size=3, padding=1, bias=False),
-            nn.BatchNorm2d(out_channels),
-            nn.GELU(),
+        self.fusion = CoastStyleMultiScaleFusion(
+            vit_dim=self.vit_channels,
+            convnext_dims=self.convnext_dims,
+            out_channels=out_channels,
+            semantic_grid_size=None,
+            detail_grid_size=(28, 28),
         )
 
         self.apply_finetune_policy(vit_unfrozen_blocks, convnext_unfrozen_stages)
 
     def apply_finetune_policy(self, vit_unfrozen_blocks: int = 2, convnext_unfrozen_stages: int = 2) -> None:
         self.requires_grad_(False)
-        self.vit_proj.requires_grad_(True)
-        self.convnext_proj.requires_grad_(True)
         self.fusion.requires_grad_(True)
         self._unfreeze_vit_tail(vit_unfrozen_blocks)
         self._unfreeze_convnext_tail(convnext_unfrozen_stages)
@@ -148,15 +296,15 @@ class DINOv3ViTConvNeXtBackbone(nn.Module):
 
     def forward(self, x: Tensor) -> dict[str, Tensor]:
         vit_map = self._vit_feature_map(x)
-        convnext_map = self._convnext_feature_map(x)
-        convnext_map = F.interpolate(convnext_map, size=vit_map.shape[-2:], mode="bilinear", align_corners=False)
-        vit_feature = self.vit_proj(vit_map)
-        convnext_feature = self.convnext_proj(convnext_map)
-        fused = self.fusion(torch.cat([vit_feature, convnext_feature], dim=1))
+        convnext_feats = self._convnext_pyramid(x)
+        fused_outputs = self.fusion(vit_map, convnext_feats)
         return {
-            "fused": fused,
-            "vit": vit_feature,
-            "convnext": convnext_feature,
+            "fused": fused_outputs["fused"],
+            "semantic": fused_outputs["semantic"],
+            "detail": fused_outputs["detail"],
+            "vit": vit_map,
+            "convnext_pyramid": convnext_feats,
+            "fusion_weights": fused_outputs["fusion_weights"],
         }
 
     def _vit_feature_map(self, x: Tensor) -> Tensor:
@@ -166,12 +314,14 @@ class DINOv3ViTConvNeXtBackbone(nn.Module):
         w = x.shape[-1] // int(getattr(self.vit, "patch_size"))
         return patch_tokens.transpose(1, 2).reshape(x.shape[0], self.vit_channels, h, w)
 
-    def _convnext_feature_map(self, x: Tensor) -> Tensor:
-        features = self.convnext.forward_features(x)
-        patch_tokens = features["x_norm_patchtokens"]
-        h = x.shape[-2] // 32
-        w = x.shape[-1] // 32
-        return patch_tokens.transpose(1, 2).reshape(x.shape[0], self.convnext_channels, h, w)
+    def _convnext_pyramid(self, x: Tensor) -> tuple[Tensor, Tensor, Tensor, Tensor]:
+        feats = []
+        y = x
+        for downsample, stage in zip(self.convnext.downsample_layers, self.convnext.stages):
+            y = downsample(y)
+            y = stage(y)
+            feats.append(y)
+        return tuple(feats[:4])
 
     def trainable_summary(self) -> TrainableBackboneSummary:
         trainable_names = [name for name, param in self.named_parameters() if param.requires_grad]