Fix model architecture: decoder_embed_dim=384, add norm layer, correct forward method

inference_engine.py

@@ -69,15 +69,16 @@ TASKS = {
 }
 
 # 模型默认参数（与训练时一致）
+# IMPORTANT: decoder_embed_dim must be divisible by 6 for 3D sincos position embedding
 DEFAULT_MODEL_ARGS = {
     'img_size': 128,
     'patch_size': 16,
     'embed_dim': 768,
     'depth': 12,
     'num_heads': 12,
-    'decoder_embed_dim': 512,
-    'decoder_depth': 8,
-    'decoder_num_heads': 16,
+    'decoder_embed_dim': 384,  # Must be divisible by 6 (384/6=64) ✓
+    'decoder_depth': 2,        # Pretrain default
+    'decoder_num_heads': 12,
     'pool': 'mean',
     'dropout': 0.5,
 }
@@ -106,7 +107,11 @@ class MultiMAE3DForDownstream(nn.Module):
         super().__init__()
         self.encoder = encoder
         self.pool = pool
-        self.num_outputs = num_outputs
+        self.num_patches_per_modality = encoder.num_patches
+        self.num_global_tokens = encoder.num_global_tokens
+        
+        # LayerNorm before head (important for checkpoint compatibility)
+        self.norm = nn.LayerNorm(embed_dim)
         
         # 预测头
         self.head = nn.Sequential(
@@ -114,30 +119,33 @@ class MultiMAE3DForDownstream(nn.Module):
             nn.Linear(embed_dim, num_outputs)
         )
         
-    def forward(self, images, observed, mc=None):
+    def forward(self, images: torch.Tensor, observed: torch.Tensor) -> torch.Tensor:
         """
         Args:
             images: [B, 4, D, H, W] - 4 modalities
             observed: [B, 4] - 0/1 mask for available modalities
-            mc: [B] - modality combination index (optional)
-        
         Returns:
             logits: [B, num_outputs]
         """
-        # Encode
-        x = self.encoder.forward_encoder(images, observed)
-        
-        # Pool
-        if self.pool == 'mean':
-            x = x.mean(dim=1)  # [B, embed_dim]
-        elif self.pool == 'cls':
-            x = x[:, 0]  # Use CLS token
+        # encode() returns [B, 1 + 4*num_patches, embed_dim]
+        encoder_out = self.encoder.encode(images, observed)
+
+        if self.pool == 'cls':
+            features = encoder_out[:, 0]  # CLS token -> [B, D]
+        elif self.pool == 'mean':
+            # Mean pool over modality tokens with masking for missing modalities
+            tokens = encoder_out[:, self.num_global_tokens:]  # [B, 4*N_p, D]
+            B, _, D = tokens.shape
+            N = self.num_patches_per_modality
+            # Build per-token mask: repeat each modality's observed flag N times
+            mask = observed.unsqueeze(-1).expand(-1, -1, N)  # [B, 4, N]
+            mask = mask.reshape(B, 4 * N).unsqueeze(-1)  # [B, 4*N, 1]
+            features = (tokens * mask).sum(dim=1) / mask.sum(dim=1).clamp(min=1.0)
         else:
-            x = x[:, 0]
-        
-        # Head
-        logits = self.head(x)  # [B, num_outputs]
-        
+            raise ValueError(f"Unknown pool type: {self.pool}")
+
+        features = self.norm(features)
+        logits = self.head(features)
         return logits