| """Encoder-Decoder Transformer for Experiment 2.c."""
|
|
|
| from __future__ import annotations
|
|
|
| import torch
|
| from torch import nn
|
|
|
| from src.models.layers import SinusoidalPositionalEncoding
|
| from src.utils.config import ModelConfig
|
|
|
|
|
| class EncoderDecoderTransformerModel(nn.Module):
|
| """Encoder consumes lookback; learned query tokens cross-attend to encoder memory → MLP head."""
|
|
|
| def __init__(
|
| self,
|
| input_dim: int,
|
| horizon: int,
|
| input_kind: str,
|
| model_config: ModelConfig,
|
| ):
|
| super().__init__()
|
| if model_config.transformer is None:
|
| raise ValueError("Transformer configuration required.")
|
|
|
| cfg = model_config.transformer
|
| num_dec = cfg.num_decoder_layers if cfg.num_decoder_layers is not None else cfg.num_layers
|
|
|
| self.input_proj = nn.Linear(input_dim, cfg.d_model)
|
| self.pos_enc = SinusoidalPositionalEncoding(cfg.d_model, dropout=cfg.dropout)
|
|
|
| encoder_layer = nn.TransformerEncoderLayer(
|
| d_model=cfg.d_model,
|
| nhead=cfg.nhead,
|
| dim_feedforward=cfg.dim_feedforward,
|
| dropout=cfg.dropout,
|
| batch_first=True,
|
| norm_first=True,
|
| )
|
| self.encoder = nn.TransformerEncoder(encoder_layer, num_layers=cfg.num_layers)
|
|
|
| decoder_layer = nn.TransformerDecoderLayer(
|
| d_model=cfg.d_model,
|
| nhead=cfg.nhead,
|
| dim_feedforward=cfg.dim_feedforward,
|
| dropout=cfg.dropout,
|
| batch_first=True,
|
| norm_first=True,
|
| )
|
| self.decoder = nn.TransformerDecoder(decoder_layer, num_layers=num_dec)
|
|
|
|
|
| self.query_tokens = nn.Parameter(torch.randn(1, horizon, cfg.d_model))
|
|
|
| self.head = nn.Sequential(
|
| nn.Linear(cfg.d_model, cfg.head_hidden_size),
|
| nn.ReLU(),
|
| nn.Dropout(cfg.dropout),
|
| nn.Linear(cfg.head_hidden_size, 1),
|
| )
|
|
|
| def forward(self, inputs: torch.Tensor) -> torch.Tensor:
|
|
|
| B = inputs.size(0)
|
| memory = self.encoder(self.pos_enc(self.input_proj(inputs)))
|
| queries = self.query_tokens.expand(B, -1, -1)
|
| out = self.decoder(tgt=queries, memory=memory)
|
| return self.head(out).squeeze(-1)
|
|
|
