Sync DINOv3 ViT ConvNeXt dual backbone (part 2)

marine_dual_dinov3_backbone.py

@@ -0,0 +1,170 @@
+"""DINOv3 ViT + ConvNeXt dual-branch backbone for marine feature models."""
+
+from __future__ import annotations
+
+import sys
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Callable
+
+import torch
+from torch import Tensor, nn
+import torch.nn.functional as F
+
+DINOV3_ROOT = Path(__file__).resolve().parent / "dinov3-main"
+if str(DINOV3_ROOT) not in sys.path:
+    sys.path.append(str(DINOV3_ROOT))
+
+from dinov3.hub import backbones as dinov3_backbones  # noqa: E402
+
+
+BatchNormTypes = (nn.BatchNorm1d, nn.BatchNorm2d, nn.BatchNorm3d, nn.SyncBatchNorm)
+
+
+@dataclass(frozen=True)
+class TrainableBackboneSummary:
+    total_tensors: int
+    trainable_tensors: int
+    trainable_parameters: int
+    trainable_prefixes: tuple[str, ...]
+
+
+def _builder(name: str) -> Callable[..., nn.Module]:
+    try:
+        return getattr(dinov3_backbones, name)
+    except AttributeError as exc:
+        raise ValueError(f"Unsupported DINOv3 backbone '{name}'") from exc
+
+
+def _set_module_trainable(module: nn.Module, trainable: bool) -> None:
+    for param in module.parameters():
+        param.requires_grad = trainable
+
+
+class DINOv3ViTConvNeXtBackbone(nn.Module):
+    """Dual-branch DINOv3 backbone with explicit partial unfreezing.
+
+    The fine-tuning policy is intentionally narrow:
+    - freeze both pretrained branches by default;
+    - unfreeze the final ``vit_unfrozen_blocks`` ViT transformer blocks;
+    - unfreeze the final ``convnext_unfrozen_stages`` ConvNeXt stages and their
+      matching downsample layers;
+    - unfreeze every BatchNorm module in the complete model.
+    """
+
+    def __init__(
+        self,
+        vit_name: str = "dinov3_vitl16",
+        convnext_name: str = "dinov3_convnext_base",
+        vit_weights: str | None = None,
+        convnext_weights: str | None = None,
+        pretrained: bool = True,
+        out_channels: int = 256,
+        vit_unfrozen_blocks: int = 2,
+        convnext_unfrozen_stages: int = 2,
+    ) -> None:
+        super().__init__()
+        self.vit_name = vit_name
+        self.convnext_name = convnext_name
+        self.vit = _builder(vit_name)(pretrained=pretrained, weights=vit_weights) if vit_weights else _builder(vit_name)(pretrained=pretrained)
+        self.convnext = (
+            _builder(convnext_name)(pretrained=pretrained, weights=convnext_weights)
+            if convnext_weights
+            else _builder(convnext_name)(pretrained=pretrained)
+        )
+        self.vit_channels = int(getattr(self.vit, "embed_dim"))
+        self.convnext_channels = int(getattr(self.convnext, "embed_dim"))
+        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.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)
+        self._unfreeze_all_batchnorm()
+
+    def _unfreeze_vit_tail(self, blocks_to_unfreeze: int) -> None:
+        blocks = getattr(self.vit, "blocks", None)
+        if blocks is None:
+            raise RuntimeError(f"{self.vit_name} does not expose a 'blocks' ModuleList")
+        for block in blocks[-blocks_to_unfreeze:]:
+            _set_module_trainable(block, True)
+
+    def _unfreeze_convnext_tail(self, stages_to_unfreeze: int) -> None:
+        stages = getattr(self.convnext, "stages", None)
+        downsample_layers = getattr(self.convnext, "downsample_layers", None)
+        if stages is None or downsample_layers is None:
+            raise RuntimeError(f"{self.convnext_name} does not expose ConvNeXt stages/downsample_layers")
+        start = max(0, len(stages) - stages_to_unfreeze)
+        for idx in range(start, len(stages)):
+            _set_module_trainable(stages[idx], True)
+            _set_module_trainable(downsample_layers[idx], True)
+
+    def _unfreeze_all_batchnorm(self) -> None:
+        for module in self.modules():
+            if isinstance(module, BatchNormTypes):
+                module.train()
+                _set_module_trainable(module, True)
+
+    def train(self, mode: bool = True) -> "DINOv3ViTConvNeXtBackbone":
+        super().train(mode)
+        if mode:
+            self._unfreeze_all_batchnorm()
+        return self
+
+    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))
+        return {
+            "fused": fused,
+            "vit": vit_feature,
+            "convnext": convnext_feature,
+        }
+
+    def _vit_feature_map(self, x: Tensor) -> Tensor:
+        features = self.vit.forward_features(x)
+        patch_tokens = features["x_norm_patchtokens"]
+        h = x.shape[-2] // int(getattr(self.vit, "patch_size"))
+        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 trainable_summary(self) -> TrainableBackboneSummary:
+        trainable_names = [name for name, param in self.named_parameters() if param.requires_grad]
+        prefixes = tuple(sorted({".".join(name.split(".")[:3]) for name in trainable_names}))
+        return TrainableBackboneSummary(
+            total_tensors=sum(1 for _ in self.parameters()),
+            trainable_tensors=len(trainable_names),
+            trainable_parameters=sum(param.numel() for param in self.parameters() if param.requires_grad),
+            trainable_prefixes=prefixes,
+        )

scripts/check_dual_backbone_policy.py

@@ -0,0 +1,65 @@
+"""Check trainable parameters for the DINOv3 ViT + ConvNeXt dual backbone."""
+
+from __future__ import annotations
+
+import argparse
+import json
+from pathlib import Path
+import sys
+
+ROOT = Path(__file__).resolve().parents[1]
+if str(ROOT) not in sys.path:
+    sys.path.append(str(ROOT))
+
+from marine_dual_dinov3_backbone import DINOv3ViTConvNeXtBackbone  # noqa: E402
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--vit-name", default="dinov3_vitl16")
+    parser.add_argument("--convnext-name", default="dinov3_convnext_base")
+    parser.add_argument("--vit-weights", default=None)
+    parser.add_argument("--convnext-weights", default=None)
+    parser.add_argument("--pretrained", action=argparse.BooleanOptionalAction, default=False)
+    parser.add_argument("--output", default=None)
+    return parser.parse_args()
+
+
+def main() -> None:
+    args = parse_args()
+    model = DINOv3ViTConvNeXtBackbone(
+        vit_name=args.vit_name,
+        convnext_name=args.convnext_name,
+        vit_weights=args.vit_weights,
+        convnext_weights=args.convnext_weights,
+        pretrained=args.pretrained,
+    )
+    trainable = [name for name, param in model.named_parameters() if param.requires_grad]
+    frozen = [name for name, param in model.named_parameters() if not param.requires_grad]
+    summary = model.trainable_summary()
+    report = {
+        "vit_name": args.vit_name,
+        "convnext_name": args.convnext_name,
+        "pretrained": args.pretrained,
+        "summary": summary.__dict__,
+        "trainable_examples": trainable[:80],
+        "frozen_examples": frozen[:40],
+        "checks": {
+            "vit_last_two_blocks_trainable": any(name.startswith("vit.blocks.22.") for name in trainable)
+            and any(name.startswith("vit.blocks.23.") for name in trainable),
+            "vit_earlier_blocks_frozen": not any(name.startswith("vit.blocks.21.") for name in trainable),
+            "convnext_last_two_stages_trainable": any(name.startswith("convnext.stages.2.") for name in trainable)
+            and any(name.startswith("convnext.stages.3.") for name in trainable),
+            "convnext_early_stages_frozen": not any(name.startswith("convnext.stages.0.") for name in trainable)
+            and not any(name.startswith("convnext.stages.1.") for name in trainable),
+            "fusion_trainable": any(name.startswith("fusion.") for name in trainable),
+        },
+    }
+    text = json.dumps(report, indent=2, ensure_ascii=False)
+    if args.output:
+        Path(args.output).write_text(text, encoding="utf-8")
+    print(text)
+
+
+if __name__ == "__main__":
+    main()