File size: 1,950 Bytes
590a501 | 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 | """Simple MLP model for factor combination."""
from __future__ import annotations
import torch
import torch.nn as nn
from torch.utils.data import DataLoader, TensorDataset
class FactorMLP(nn.Module):
def __init__(self, n_features: int, hidden: int = 64):
super().__init__()
self.net = nn.Sequential(
nn.Linear(n_features, hidden),
nn.ReLU(),
nn.Dropout(0.2),
nn.Linear(hidden, hidden // 2),
nn.ReLU(),
nn.Linear(hidden // 2, 1),
)
def forward(self, x):
return self.net(x).squeeze(-1)
def train_mlp(X_train, y_train, X_valid, y_valid, epochs: int = 20, lr: float = 1e-3):
device = "cuda" if torch.cuda.is_available() else "cpu"
model = FactorMLP(X_train.shape[1]).to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=lr)
loss_fn = nn.MSELoss()
train_loader = DataLoader(
TensorDataset(torch.tensor(X_train, dtype=torch.float32), torch.tensor(y_train, dtype=torch.float32)),
batch_size=4096,
shuffle=True,
)
X_valid_t = torch.tensor(X_valid, dtype=torch.float32, device=device)
y_valid_t = torch.tensor(y_valid, dtype=torch.float32, device=device)
best_loss, best_state = float("inf"), None
for _ in range(epochs):
model.train()
for xb, yb in train_loader:
xb, yb = xb.to(device), yb.to(device)
optimizer.zero_grad()
loss = loss_fn(model(xb), yb)
loss.backward()
optimizer.step()
model.eval()
with torch.no_grad():
valid_loss = loss_fn(model(X_valid_t), y_valid_t).item()
if valid_loss < best_loss:
best_loss = valid_loss
best_state = {k: v.cpu().clone() for k, v in model.state_dict().items()}
if best_state:
model.load_state_dict(best_state)
print(f"MLP valid MSE: {best_loss:.6f}")
return model
|