Update model.py

updated model module

model.py

@@ -2,11 +2,21 @@
 StableResNet Model for Biomass Prediction
 A numerically stable ResNet architecture for regression tasks
 
+Author: najahpokkiri
+Date: 2025-05-17
+"""
+"""
+StableResNet Model Architecture
+
+This module defines the StableResNet architecture used for biomass prediction.
+The model is designed for numerical stability with batch normalization and residual connections.
+
 Author: najahpokkiri
 Date: 2025-05-17
 """
 import torch
 import torch.nn as nn
+import torch.nn.functional as F
 
 class StableResNet(nn.Module):
     """Numerically stable ResNet for biomass regression"""
@@ -33,20 +43,27 @@ class StableResNet(nn.Module):
         self._init_weights()
     
     def _make_simple_resblock(self, in_dim, out_dim):
-        return nn.Sequential(
-            nn.Linear(in_dim, out_dim),
-            nn.BatchNorm1d(out_dim),
-            nn.ReLU(),
-            nn.Linear(out_dim, out_dim),
-            nn.BatchNorm1d(out_dim),
-            nn.ReLU()
-        ) if in_dim == out_dim else nn.Sequential(
-            nn.Linear(in_dim, out_dim),
-            nn.BatchNorm1d(out_dim),
-            nn.ReLU(),
-        )
+        """Create a simple residual block or downsampling block"""
+        if in_dim == out_dim:
+            # Residual block
+            return nn.Sequential(
+                nn.Linear(in_dim, out_dim),
+                nn.BatchNorm1d(out_dim),
+                nn.ReLU(),
+                nn.Linear(out_dim, out_dim),
+                nn.BatchNorm1d(out_dim),
+                nn.ReLU()
+            )
+        else:
+            # Downsampling block
+            return nn.Sequential(
+                nn.Linear(in_dim, out_dim),
+                nn.BatchNorm1d(out_dim),
+                nn.ReLU(),
+            )
     
     def _init_weights(self):
+        """Initialize weights for better convergence"""
         for m in self.modules():
             if isinstance(m, nn.Linear):
                 nn.init.kaiming_normal_(m.weight, mode='fan_in', nonlinearity='relu')
@@ -54,14 +71,35 @@ class StableResNet(nn.Module):
                     nn.init.zeros_(m.bias)
     
     def forward(self, x):
+        """Forward pass through the network"""
         x = self.input_proj(x)
         
+        # First residual block
         identity = x
         out = self.layer1(x)
         x = out + identity
         
+        # Remaining blocks
         x = self.layer2(x)
         x = self.layer3(x)
         
+        # Regression output
         x = self.regressor(x)
-        return x.squeeze()
+        return x.squeeze()
+
+def get_model_info():
+    """Return information about the model architecture"""
+    return {
+        'name': 'StableResNet',
+        'description': 'Numerically stable ResNet for biomass regression',
+        'parameters': {
+            'n_features': 'Number of input features',
+            'dropout': 'Dropout rate (default: 0.2)'
+        },
+        'architecture': [
+            'Input projection with layer normalization',
+            'Residual blocks with batch normalization',
+            'Downsampling blocks',
+            'Regression head'
+        ]
+    }