File size: 4,328 Bytes
6cb1024 | 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 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 | import json
import torch
from torch import nn
class Model(nn.Module):
"""LSTM + MLP regressor that predicts the next step's numeric features.
Each timestep is a vector of ``numeric_features`` normalized values
followed by an integer ticker id, which is embedded and concatenated
to the numeric features before the LSTM.
"""
def __init__(
self,
num_tickers,
numeric_features=5,
embedding_dim=16,
output_size=5,
):
super().__init__()
self.numeric_features = numeric_features
self.ticker_embedding = nn.Embedding(
num_embeddings=num_tickers + 1,
embedding_dim=embedding_dim,
)
self.lstm = nn.LSTM(
input_size=numeric_features + embedding_dim,
hidden_size=128,
num_layers=3,
batch_first=True,
)
self.mlp = nn.Sequential(
nn.Linear(128, 128),
nn.GELU(),
nn.Linear(128, 256),
nn.GELU(),
nn.Dropout(0.2),
nn.Linear(256, 128),
nn.GELU(),
nn.Linear(128, output_size),
)
def forward(self, x):
"""Predict the next step.
Args:
x: Tensor of shape (batch, sequence, numeric_features + 1) where
the last channel holds the ticker id.
Returns:
Tensor of shape (batch, output_size).
"""
numeric = x[..., : self.numeric_features]
ticker_ids = x[..., self.numeric_features].long()
ticker_vectors = self.ticker_embedding(ticker_ids)
x = torch.cat([numeric, ticker_vectors], dim=-1)
output, _ = self.lstm(x)
return self.mlp(output[:, -1, :])
class Normalizer(object):
"""Applies the per-ticker mean/std normalization stored in a config file."""
def __init__(self, config_path: str | None = None):
"""Load the normalization config.
Args:
config_path: Path to the JSON config. Defaults to `./config.json`.
"""
if config_path is None:
config_path = "./config.json"
with open(config_path) as f:
self.config = json.load(f)
self.ticker_ids = self.config["ticker_ids"]
self.mean = torch.tensor(self.config["mean"], dtype=torch.float)
self.std = torch.tensor(self.config["std"], dtype=torch.float)
def __call__(self, x, ticker_name) -> torch.Tensor:
"""Normalize ``x`` with the statistics of ``ticker_name``."""
return self.normalize(x, ticker_name)
def normalize(self, x, ticker_name: str) -> torch.Tensor:
"""Return ``(x - mean) / std`` for the given ticker.
Args:
x: Array-like of shape (..., num_features). Volume is expected to
already be log1p-transformed, as in training.
ticker_name: Ticker symbol present in the config's ``ticker_ids``.
Raises:
KeyError: If the ticker is not in the config.
"""
ticker_id = self.ticker_ids[ticker_name]
x = torch.as_tensor(x, dtype=torch.float)
return (x - self.mean[ticker_id]) / self.std[ticker_id]
def load_model(
weights_path: str = "stock_predictor/model.pth",
config_path: str | None = None,
device: str | torch.device | None = None,
) -> Model:
"""Build a ``Model`` from the config and load trained weights into it.
Args:
weights_path: Path to the saved ``state_dict``.
config_path: Path to the JSON config. Defaults to ``./config.json``.
device: Device to load onto. Defaults to CUDA if available, else CPU.
Returns:
The model on ``device`` in evaluation mode.
"""
if config_path is None:
config_path = "./config.json"
if device is None:
device = "cuda" if torch.cuda.is_available() else "cpu"
device = torch.device(device)
with open(config_path) as f:
config = json.load(f)
model = Model(
num_tickers=len(config["ticker_ids"]),
numeric_features=len(config["mean"][0]),
output_size=len(config["mean"][0]),
)
state_dict = torch.load(weights_path, map_location=device, weights_only=True)
model.load_state_dict(state_dict)
model.to(device)
model.eval()
return model
|