Simplify the head for segmentation

Author: srmsoumya

configs/segment_chesapeake.yaml

@@ -11,12 +11,7 @@ data:
   platform: naip
 model:
   num_classes: 7
-  feature_maps:
-    - 5
-    - 11
-    - 15
-    - 23
-  ckpt_path: checkpoints/v1.5.0-no-mrl-dinov2/mae_v1.5.0_epoch-05_val-loss-0.1734.ckpt
+  ckpt_path: checkpoints/clay_v1.5.ckpt
   lr: 1e-5
   wd: 0.05
   b1: 0.9
@@ -28,7 +23,7 @@ trainer:
   num_nodes: 1
   precision: bf16-mixed
   log_every_n_steps: 5
-  max_epochs: 10
+  max_epochs: 100
   accumulate_grad_batches: 1
   default_root_dir: checkpoints/segment
   fast_dev_run: False

finetune/segment/chesapeake_model.py

@@ -28,7 +28,6 @@ class ChesapeakeSegmentor(L.LightningModule):
     def __init__(  # # noqa: PLR0913
         self,
         num_classes,
-        feature_maps,
         ckpt_path,
         lr,
         wd,
@@ -39,7 +38,6 @@ def __init__(  # # noqa: PLR0913
         self.save_hyperparameters()  # Save hyperparameters for checkpointing
         self.model = Segmentor(
             num_classes=num_classes,
-            feature_maps=feature_maps,
             ckpt_path=ckpt_path,
         )
 
@@ -101,7 +99,7 @@ def configure_optimizers(self):
             optimizer,
             T_0=100,
             T_mult=1,
-            eta_min=self.hparams.lr * 10,
+            eta_min=self.hparams.lr * 100,
             last_epoch=-1,
         )
         return {

finetune/segment/factory.py

@@ -10,6 +10,7 @@
 import re
 
 import torch
+import torch.nn.functional as F
 from einops import rearrange, repeat
 from torch import nn
 
@@ -37,7 +38,6 @@ def __init__(  # noqa: PLR0913
         heads,
         dim_head,
         mlp_ratio,
-        feature_maps,
         ckpt_path=None,
     ):
         super().__init__(
@@ -50,27 +50,6 @@ def __init__(  # noqa: PLR0913
             dim_head,
             mlp_ratio,
         )
-        self.feature_maps = feature_maps
-
-        # Define Feature Pyramid Network (FPN) layers
-        self.fpn1 = nn.Sequential(
-            nn.ConvTranspose2d(dim, dim, kernel_size=2, stride=2),
-            nn.BatchNorm2d(dim),
-            nn.GELU(),
-            nn.ConvTranspose2d(dim, dim, kernel_size=2, stride=2),
-        )
-
-        self.fpn2 = nn.Sequential(
-            nn.ConvTranspose2d(dim, dim, kernel_size=2, stride=2),
-        )
-
-        self.fpn3 = nn.Identity()
-
-        self.fpn4 = nn.Sequential(
-            nn.MaxPool2d(kernel_size=2, stride=2),
-        )
-
-        self.fpn5 = nn.Identity()
 
         # Set device
         self.device = (
@@ -143,25 +122,10 @@ def forward(self, datacube):
         cls_tokens = repeat(self.cls_token, "1 1 D -> B 1 D", B=B)  # [B 1 D]
         patches = torch.cat((cls_tokens, patches), dim=1)  # [B (1 + L) D]
 
-        features = []
-        for idx, (attn, ff) in enumerate(self.transformer.layers):
-            patches = attn(patches) + patches
-            patches = ff(patches) + patches
-            if idx in self.feature_maps:
-                _cube = rearrange(
-                    patches[:, 1:, :], "B (H W) D -> B D H W", H=H // 8, W=W // 8
-                )
-                features.append(_cube)
-        patches = self.transformer.norm(patches)
-        _cube = rearrange(patches[:, 1:, :], "B (H W) D -> B D H W", H=H // 8, W=W // 8)
-        features.append(_cube)
-
-        # Apply FPN layers
-        ops = [self.fpn1, self.fpn2, self.fpn3, self.fpn4, self.fpn5]
-        for i in range(len(features)):
-            features[i] = ops[i](features[i])
+        patches = self.transformer(patches)
+        patches = patches[:, 1:, :]  # [B L D]
 
-        return features
+        return patches
 
 
 class Segmentor(nn.Module):
@@ -175,7 +139,7 @@ class Segmentor(nn.Module):
         ckpt_path (str): Path to the checkpoint file.
     """
 
-    def __init__(self, num_classes, feature_maps, ckpt_path):
+    def __init__(self, num_classes, ckpt_path):
         super().__init__()
         # Default values are for the clay mae base model.
         self.encoder = SegmentEncoder(
@@ -187,14 +151,26 @@ def __init__(self, num_classes, feature_maps, ckpt_path):
             heads=16,
             dim_head=64,
             mlp_ratio=4.0,
-            feature_maps=feature_maps,
             ckpt_path=ckpt_path,
         )
-        self.upsamples = [nn.Upsample(scale_factor=2**i) for i in range(4)] + [
-            nn.Upsample(scale_factor=4),
-        ]
-        self.fusion = nn.Conv2d(self.encoder.dim * 5, self.encoder.dim, kernel_size=1)
-        self.seg_head = nn.Conv2d(self.encoder.dim, num_classes, kernel_size=1)
+
+        # Freeze the encoder parameters
+        for param in self.encoder.parameters():
+            param.requires_grad = False
+
+         # Define layers after the encoder
+        D = self.encoder.dim  # embedding dimension
+        hidden_dim = 512
+        C_out = 64
+        r = self.encoder.patch_size  # upscale factor (patch_size)
+
+        self.conv1 = nn.Conv2d(D, hidden_dim, kernel_size=3, padding=1)
+        self.bn1 = nn.BatchNorm2d(hidden_dim)
+        self.conv2 = nn.Conv2d(hidden_dim, hidden_dim, kernel_size=3, padding=1)
+        self.bn2 = nn.BatchNorm2d(hidden_dim)
+        self.conv_ps = nn.Conv2d(hidden_dim, C_out * r * r, kernel_size=3, padding=1)
+        self.pixel_shuffle = nn.PixelShuffle(upscale_factor=r)
+        self.conv_out = nn.Conv2d(C_out, num_classes, kernel_size=3, padding=1)
 
     def forward(self, datacube):
         """
@@ -207,12 +183,26 @@ def forward(self, datacube):
         Returns:
             torch.Tensor: The segmentation logits.
         """
-        features = self.encoder(datacube)
-        for i in range(len(features)):
-            features[i] = self.upsamples[i](features[i])
+        cube = datacube["pixels"]  # [B C H_in W_in]
+        B, C, H_in, W_in = cube.shape
+
+        # Get embeddings from the encoder
+        patches = self.encoder(datacube)  # [B, L, D]
+
+        # Reshape embeddings to [B, D, H', W']
+        H_patches = H_in // self.encoder.patch_size
+        W_patches = W_in // self.encoder.patch_size
+        x = rearrange(patches, "B (H W) D -> B D H W", H=H_patches, W=W_patches)
+
+        # Pass through convolutional layers
+        x = F.relu(self.bn1(self.conv1(x)))
+        x = F.relu(self.bn2(self.conv2(x)))
+        x = self.conv_ps(x)  # [B, C_out * r^2, H', W']
+
+        # Upsample using PixelShuffle
+        x = self.pixel_shuffle(x)  # [B, C_out, H_in, W_in]
 
-        fused = torch.cat(features, dim=1)
-        fused = self.fusion(fused)
+        # Final convolution to get desired output channels
+        x = self.conv_out(x)  # [B, num_outputs, H_in, W_in]
 
-        logits = self.seg_head(fused)
-        return logits
+        return x