| """ |
| StableResNet Model for Biomass Prediction |
| A numerically stable ResNet architecture for regression tasks |
| |
| Author: najahpokkiri |
| Date: 2025-05-17 |
| """ |
| import torch |
| import torch.nn as nn |
|
|
| class StableResNet(nn.Module): |
| """Numerically stable ResNet for biomass regression""" |
| def __init__(self, n_features, dropout=0.2): |
| super().__init__() |
| |
| self.input_proj = nn.Sequential( |
| nn.Linear(n_features, 256), |
| nn.LayerNorm(256), |
| nn.ReLU(), |
| nn.Dropout(dropout) |
| ) |
| |
| self.layer1 = self._make_simple_resblock(256, 256) |
| self.layer2 = self._make_simple_resblock(256, 128) |
| self.layer3 = self._make_simple_resblock(128, 64) |
| |
| self.regressor = nn.Sequential( |
| nn.Linear(64, 32), |
| nn.ReLU(), |
| nn.Linear(32, 1) |
| ) |
| |
| 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(), |
| ) |
| |
| def _init_weights(self): |
| for m in self.modules(): |
| if isinstance(m, nn.Linear): |
| nn.init.kaiming_normal_(m.weight, mode='fan_in', nonlinearity='relu') |
| if m.bias is not None: |
| nn.init.zeros_(m.bias) |
| |
| def forward(self, x): |
| x = self.input_proj(x) |
| |
| identity = x |
| out = self.layer1(x) |
| x = out + identity |
| |
| x = self.layer2(x) |
| x = self.layer3(x) |
| |
| x = self.regressor(x) |
| return x.squeeze() |
