baseline.py

import torch
from torch import nn
import torch.nn.functional as F
window_len = 40

# ==========================================
# SHARED MODULES
# ==========================================

class InputProjector(nn.Module):
    """
    The initial projection block specified:
    Linear -> LRelu -> Linear -> LRelu -> Linear
    """
    def __init__(self, input_dim=17, hidden_dim=32, out_dim=64):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(input_dim, hidden_dim),
            nn.LeakyReLU(),
            nn.Linear(hidden_dim, out_dim),
            nn.LeakyReLU(),
            nn.Linear(out_dim, out_dim) # Final projection to d_model
        )

    def forward(self, x):
        return self.net(x)

class OutputDec(nn.Module):
    """
    The final regression head:
    Linear -> Lrelu -> Linear -> RUL (dim 1)
    """
    def __init__(self, input_dim=64, hidden_dim=32):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(input_dim, hidden_dim),
            nn.LeakyReLU(),
            nn.Linear(hidden_dim, 1) # Output is scalar RUL
        )

    def forward(self, x):
        return self.net(x)

class AttentionPooling(nn.Module):
    """
    Latent obtained by sum with attention weights.
    """
    def __init__(self, dim):
        super().__init__()
        self.att = nn.Linear(dim, 1)

    def forward(self, x):
        # x shape: [Batch, Window, Dim]
        # weights: [Batch, Window, 1]
        w = torch.softmax(self.att(x), dim=1)
        # Weighted sum: [Batch, Dim]
        return (x * w).sum(dim=1)

# ==========================================
# 1. LSTM BASELINE
# ==========================================

class LSTMBaseline(nn.Module):
    def __init__(self, input_dim=17, d_model=64):
        super().__init__()
        
        # 1. Input Processing
        self.proj = InputProjector(input_dim=input_dim, out_dim=d_model)
        
        # 2. LSTM Core
        # "2 cells and zero initial hidden and cell state"
        # PyTorch initializes h0, c0 to zeros by default if not provided
        self.lstm = nn.LSTM(
            input_size=d_model, 
            hidden_size=d_model, 
            num_layers=2, 
            batch_first=True
        )
        
        # 3. Pooling
        self.pool = AttentionPooling(d_model)
        
        # 4. Output Head
        self.head = OutputDec(input_dim=d_model)

    def forward(self, x):
        # x: [Batch, Window, 17]
        
        # Project: [Batch, Window, d_model]
        x = self.proj(x)
        
        # LSTM: [Batch, Window, d_model]
        # We ignore the hidden states (h_n, c_n) and take the sequence output
        x, _ = self.lstm(x)
        
        # Attention Pooling: [Batch, d_model]
        latent = self.pool(x)
        
        # Regression: [Batch, 1]
        rul = self.head(latent)
        return rul

# ==========================================
# 2. TRANSFORMER BASELINE
# ==========================================

class TransformerBaseline(nn.Module):
    def __init__(self, input_dim=17, d_model=64, nhead=4, dim_feedforward=256, proj_hidden_dim=32, num_layers=3):
        super().__init__()
        
        # 1. Input Processing
        self.proj = InputProjector(input_dim=input_dim, hidden_dim=proj_hidden_dim, out_dim=d_model)
        
        # Positional Encoding (Learnable)
        self.pos = nn.Parameter(torch.randn(1, window_len, d_model))
        
        # 2. Transformer Core
        # "3 encoder heads and one decoder heads with gelu loss"
        # We use pytorch's prebuilt transformer module
        encoder_layer = nn.TransformerEncoderLayer(
            d_model=d_model,
            nhead=nhead,            # 3 encoder heads (or 4 to match d_model div)
            dim_feedforward=dim_feedforward,
            dropout=0.1,
            activation='gelu',      # gelu activation
            batch_first=True
        )
        self.trf_encoder = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)
        
        # 3. Pooling
        self.pool = AttentionPooling(d_model)
        
        # 4. Output Head ("one decoder head")
        self.head = OutputDec(input_dim=d_model)

    def forward(self, x):
        # x: [Batch, Window, 17]
        
        # Project + Positional Encoding
        x = self.proj(x) + self.pos
        
        # Transformer Encoder
        x = self.trf_encoder(x)
        
        # Attention Pooling
        latent = self.pool(x)
        
        # Regression
        rul = self.head(latent)
        return rul