File size: 1,080 Bytes
85a1a1d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 | import torch
import torch.nn as nn
class TinyTransformerForecaster(nn.Module):
"""
Small Transformer encoder for time-series (CPU-friendly).
Input: (B, T, F=4)
Output: (B, 1) predicted log-return
"""
def __init__(self, feature_dim=4, d_model=64, nhead=4, num_layers=2, dropout=0.1):
super().__init__()
self.in_proj = nn.Linear(feature_dim, d_model)
encoder_layer = nn.TransformerEncoderLayer(
d_model=d_model,
nhead=nhead,
dim_feedforward=4 * d_model,
dropout=dropout,
batch_first=True,
activation="gelu",
norm_first=True,
)
self.encoder = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)
self.head = nn.Sequential(
nn.LayerNorm(d_model),
nn.Linear(d_model, d_model // 2),
nn.GELU(),
nn.Linear(d_model // 2, 1),
)
def forward(self, x):
h = self.in_proj(x)
h = self.encoder(h)
last = h[:, -1, :]
return self.head(last)
|