wehub-resource-sync
684 lines
20 KiB
Python
684 lines
20 KiB
Python
import copy
|
|
from collections import defaultdict
|
|
from typing import Literal, cast
|
|
|
|
import numpy as np
|
|
import pandas as pd
|
|
import polars as pl
|
|
from sklearn.metrics import mean_squared_error # type: ignore
|
|
import torch
|
|
import torch.nn as nn
|
|
import torch.optim as optim
|
|
|
|
from vnpy.alpha import (
|
|
AlphaDataset,
|
|
AlphaModel,
|
|
Segment,
|
|
logger
|
|
)
|
|
|
|
|
|
|
|
class MlpModel(AlphaModel):
|
|
"""
|
|
Multi-Layer Perceptron Model
|
|
|
|
Alpha factor prediction model implemented using multi-layer perceptron, with main features including:
|
|
1. Building and training multi-layer perceptron neural networks
|
|
2. Predicting Alpha factor values
|
|
3. Model evaluation and feature importance analysis
|
|
4. Support for early stopping and overfitting prevention
|
|
5. Support for MSE loss function
|
|
6. Optional Adam or SGD optimizer
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
input_size: int,
|
|
hidden_sizes: tuple[int] = (256,),
|
|
lr: float = 0.001,
|
|
n_epochs: int = 300,
|
|
batch_size: int = 2000,
|
|
early_stop_rounds: int = 50,
|
|
eval_steps: int = 20,
|
|
optimizer: Literal["sgd", "adam"] = "adam",
|
|
weight_decay: float = 0.0,
|
|
device: str = "cpu",
|
|
seed: int | None = None
|
|
) -> None:
|
|
"""
|
|
Initialize MLP model
|
|
|
|
Parameters
|
|
----------
|
|
input_size : int, default 360
|
|
Input feature dimension
|
|
hidden_sizes : tuple[int], default (256,)
|
|
Number of neurons in hidden layers
|
|
lr : float, default 0.001
|
|
Learning rate
|
|
n_epochs : int, default 300
|
|
Maximum training steps
|
|
batch_size : int, default 2000
|
|
Number of samples per batch
|
|
early_stop_rounds : int, default 50
|
|
Early stopping rounds, training stops if validation loss doesn't improve within these rounds
|
|
eval_steps : int, default 20
|
|
Evaluate model every this many steps
|
|
optimizer : Literal["sgd", "adam"], default "adam"
|
|
Optimizer type, options are "sgd" or "adam"
|
|
weight_decay : float, default 0.0
|
|
L2 regularization coefficient
|
|
seed : Optional[int], optional
|
|
Random seed for reproducibility
|
|
device : str, default "cpu"
|
|
Training device
|
|
"""
|
|
# Save model hyperparameters
|
|
self.input_size: int = input_size
|
|
self.hidden_sizes: tuple[int] = hidden_sizes
|
|
self.lr: float = lr
|
|
self.n_epochs: int = n_epochs
|
|
self.batch_size: int = batch_size
|
|
self.early_stop_rounds: int = early_stop_rounds
|
|
self.eval_steps: int = eval_steps
|
|
self.device: str = device
|
|
self.fitted: bool = False
|
|
self.feature_names: list[str] = []
|
|
self.best_step: int | None = None
|
|
|
|
# Set random seed for reproducibility
|
|
if seed is not None:
|
|
np.random.seed(seed)
|
|
torch.manual_seed(seed)
|
|
|
|
# Set loss function type
|
|
self._scorer = mean_squared_error
|
|
|
|
# Initialize model
|
|
self.model: nn.Module = MlpNetwork(
|
|
input_size=input_size,
|
|
hidden_sizes=hidden_sizes,
|
|
)
|
|
|
|
# Move model to specified device
|
|
self.model = self.model.to(device)
|
|
|
|
# Set optimizer
|
|
optimizer_name = optimizer.lower()
|
|
if optimizer_name == "adam":
|
|
self.optimizer: optim.Optimizer = optim.Adam(
|
|
self.model.parameters(),
|
|
lr=lr,
|
|
weight_decay=weight_decay
|
|
)
|
|
elif optimizer_name == "sgd":
|
|
self.optimizer = optim.SGD(
|
|
self.model.parameters(),
|
|
lr=lr,
|
|
weight_decay=weight_decay
|
|
)
|
|
else:
|
|
raise NotImplementedError(f"optimizer {optimizer} is not supported!")
|
|
|
|
# Set learning rate scheduler
|
|
self.scheduler: optim.lr_scheduler.ReduceLROnPlateau = optim.lr_scheduler.ReduceLROnPlateau(
|
|
self.optimizer,
|
|
mode="min",
|
|
factor=0.5,
|
|
patience=10,
|
|
threshold=0.0001,
|
|
threshold_mode="rel",
|
|
cooldown=0,
|
|
min_lr=0.00001,
|
|
eps=1e-08,
|
|
)
|
|
|
|
def fit(
|
|
self,
|
|
dataset: AlphaDataset,
|
|
evaluation_results: dict | None = None,
|
|
) -> None:
|
|
"""
|
|
Train the multi-layer perceptron model
|
|
|
|
Trains the MLP model using the given dataset, with main steps including:
|
|
1. Preparing training and validation data
|
|
2. Iteratively training for multiple steps
|
|
3. Evaluating model performance at fixed intervals
|
|
4. Implementing early stopping to prevent overfitting
|
|
|
|
Parameters
|
|
----------
|
|
dataset : AlphaDataset
|
|
Dataset object containing training data
|
|
evaluation_results : dict
|
|
Dictionary for storing evaluation metrics during training
|
|
"""
|
|
# Initialize a new dictionary if evaluation_results is None
|
|
if evaluation_results is None:
|
|
evaluation_results = {}
|
|
|
|
# Dictionary to store training and validation data
|
|
train_valid_data: dict[str, dict] = defaultdict(dict)
|
|
|
|
# Process training and validation sets separately
|
|
for segment in [Segment.TRAIN, Segment.VALID]:
|
|
# Get learning data and sort by time and trading code
|
|
df: pl.DataFrame = dataset.fetch_learn(segment)
|
|
df = df.sort(["datetime", "vt_symbol"])
|
|
|
|
# Extract features and labels
|
|
features = df.select(df.columns[2: -1]).to_numpy()
|
|
labels = np.array(df["label"])
|
|
|
|
# Store feature and label data
|
|
train_valid_data["x"][segment] = torch.from_numpy(features).float().to(self.device)
|
|
train_valid_data["y"][segment] = torch.from_numpy(labels).float().to(self.device)
|
|
|
|
# Initialize evaluation results list
|
|
evaluation_results[segment] = []
|
|
|
|
# Get feature names
|
|
df = dataset.fetch_learn(Segment.TRAIN)
|
|
self.feature_names = df.columns[2:-1]
|
|
|
|
# Initialize training state
|
|
early_stop_count: int = 0 # Number of steps without performance improvement
|
|
train_loss: float = 0 # Current training loss
|
|
best_valid_score: float = np.inf # Best validation loss
|
|
best_params = None # Best model parameters
|
|
|
|
train_samples: int = train_valid_data["y"][Segment.TRAIN].shape[0]
|
|
|
|
# Iterate through training steps
|
|
for step in range(1, self.n_epochs + 1):
|
|
# Check if early stopping condition is met
|
|
if early_stop_count >= self.early_stop_rounds:
|
|
logger.info("达到早停条件,训练结束")
|
|
break
|
|
|
|
# Train one batch
|
|
batch_loss = self._train_step(train_valid_data, train_samples)
|
|
train_loss += batch_loss
|
|
|
|
# Periodically evaluate the model
|
|
if step % self.eval_steps == 0 or step == self.n_epochs:
|
|
early_stop_count, best_valid_score, best_params = self._evaluate_step(
|
|
train_valid_data,
|
|
evaluation_results,
|
|
step,
|
|
train_loss,
|
|
early_stop_count,
|
|
best_valid_score
|
|
)
|
|
train_loss = 0
|
|
|
|
# Mark model as trained
|
|
self.fitted = True
|
|
|
|
# Load best model parameters
|
|
if best_params:
|
|
self.model.load_state_dict(best_params)
|
|
|
|
def _train_step(
|
|
self,
|
|
train_valid_data: dict[str, dict[Segment, torch.Tensor]],
|
|
train_samples: int
|
|
) -> float:
|
|
"""
|
|
Execute one training step
|
|
|
|
Parameters
|
|
----------
|
|
train_valid_data : dict
|
|
Training and validation data
|
|
train_samples : int
|
|
Number of training samples
|
|
|
|
Returns
|
|
-------
|
|
float
|
|
Current batch loss value
|
|
"""
|
|
batch_loss = AverageMeter()
|
|
self.model.train()
|
|
self.optimizer.zero_grad()
|
|
|
|
# Randomly select batch data
|
|
batch_indices = np.random.choice(train_samples, self.batch_size)
|
|
batch_features = train_valid_data["x"][Segment.TRAIN][batch_indices]
|
|
batch_labels = train_valid_data["y"][Segment.TRAIN][batch_indices]
|
|
|
|
# Forward and backward propagation
|
|
predictions = self.model(batch_features)
|
|
cur_loss = self._loss_fn(predictions, batch_labels)
|
|
cur_loss.backward()
|
|
|
|
# Update model parameters
|
|
self.optimizer.step()
|
|
batch_loss.update(cur_loss.item())
|
|
|
|
return batch_loss.val
|
|
|
|
def _evaluate_step(
|
|
self,
|
|
train_valid_data: dict[str, dict[Segment, torch.Tensor]],
|
|
evaluation_results: dict[Segment, list[float]],
|
|
step: int,
|
|
train_loss: float,
|
|
early_stop_count: int,
|
|
best_valid_score: float
|
|
) -> tuple[int, float, dict[str, torch.Tensor] | None]:
|
|
"""
|
|
Evaluate current model performance
|
|
|
|
Parameters
|
|
----------
|
|
train_valid_data : dict
|
|
Training and validation data
|
|
evaluation_results : dict
|
|
Evaluation results record
|
|
step : int
|
|
Current training step
|
|
train_loss : float
|
|
Current training loss
|
|
early_stop_count : int
|
|
Count of steps without improvement
|
|
best_valid_score : float
|
|
Best validation loss
|
|
|
|
Returns
|
|
-------
|
|
tuple[int, float, dict] | None
|
|
Returns updated early stop count, best validation loss, and best model parameters
|
|
"""
|
|
early_stop_count += 1
|
|
train_loss /= self.eval_steps
|
|
|
|
# Evaluate model on validation set
|
|
with torch.no_grad():
|
|
self.model.eval()
|
|
|
|
data: torch.Tensor = train_valid_data["x"][Segment.VALID]
|
|
pred: torch.Tensor = cast(torch.Tensor, self._predict_batch(data, return_cpu=False))
|
|
valid_loss = self._loss_fn(pred, train_valid_data["y"][Segment.VALID])
|
|
|
|
loss_val = valid_loss.item()
|
|
|
|
# Record evaluation results
|
|
logger.info(f"[Step {step}]: train_loss {train_loss:.6f}, valid_loss {loss_val:.6f}")
|
|
evaluation_results[Segment.TRAIN].append(train_loss)
|
|
evaluation_results[Segment.VALID].append(loss_val)
|
|
|
|
# Update best model if validation performance improves
|
|
best_params = None
|
|
if loss_val < best_valid_score:
|
|
logger.info(f"\t验证集损失从 {best_valid_score:.6f} 降低到 {loss_val:.6f}")
|
|
best_valid_score = loss_val
|
|
self.best_step = step
|
|
early_stop_count = 0
|
|
best_params = copy.deepcopy(self.model.state_dict())
|
|
|
|
# Update learning rate
|
|
if self.scheduler is not None:
|
|
self.scheduler.step(metrics=valid_loss, epoch=step)
|
|
|
|
return early_stop_count, best_valid_score, best_params
|
|
|
|
def _loss_fn(self, pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
|
|
"""
|
|
Calculate loss value
|
|
|
|
Parameters
|
|
----------
|
|
pred : torch.Tensor
|
|
Model predictions
|
|
target : torch.Tensor
|
|
Target true values
|
|
|
|
Returns
|
|
-------
|
|
torch.Tensor
|
|
Calculated loss value
|
|
"""
|
|
pred, target = pred.reshape(-1), target.reshape(-1)
|
|
loss: torch.Tensor = nn.MSELoss()(pred, target)
|
|
return loss
|
|
|
|
def _predict_batch(self, data: torch.Tensor, return_cpu: bool = True) -> np.ndarray | torch.Tensor:
|
|
"""
|
|
Neural network prediction function
|
|
|
|
Parameters
|
|
----------
|
|
data : torch.Tensor
|
|
Input data
|
|
return_cpu : bool, default True
|
|
Whether to return CPU tensor
|
|
step : Optional[int], optional
|
|
Current training step
|
|
|
|
Returns
|
|
-------
|
|
np.ndarray | torch.Tensor
|
|
Model prediction results
|
|
"""
|
|
data = data.to(self.device)
|
|
|
|
predictions: list[torch.Tensor] = []
|
|
|
|
self.model.eval()
|
|
|
|
with torch.no_grad():
|
|
batch_size: int = 8096
|
|
for i in range(0, len(data), batch_size):
|
|
x: torch.Tensor = data[i: i + batch_size]
|
|
predictions.append(self.model(x.to(self.device)).detach().reshape(-1))
|
|
|
|
if return_cpu:
|
|
return np.concatenate([pr.cpu().numpy() for pr in predictions])
|
|
else:
|
|
return torch.cat(predictions, dim=0)
|
|
|
|
def predict(self, dataset: AlphaDataset, segment: Segment) -> np.ndarray:
|
|
"""
|
|
Model prediction interface
|
|
|
|
Parameters
|
|
----------
|
|
dataset : AlphaDataset
|
|
Prediction dataset
|
|
segment : Segment
|
|
Dataset segment
|
|
|
|
Returns
|
|
-------
|
|
np.ndarray
|
|
Prediction result array
|
|
"""
|
|
if not self.fitted:
|
|
raise ValueError("Model has not been trained yet!")
|
|
|
|
df: pl.DataFrame = dataset.fetch_infer(segment)
|
|
df = df.sort(["datetime", "vt_symbol"])
|
|
|
|
data: np.ndarray = df.select(df.columns[2: -1]).to_numpy()
|
|
|
|
return cast(np.ndarray, self._predict_batch(torch.Tensor(data)))
|
|
|
|
def _check_tensor_nan(self, tensor: torch.Tensor, name: str) -> None:
|
|
"""
|
|
Check if tensor contains NaN values
|
|
|
|
Parameters
|
|
----------
|
|
tensor : torch.Tensor
|
|
Tensor to check
|
|
name : str
|
|
Tensor name
|
|
|
|
Returns
|
|
-------
|
|
None
|
|
"""
|
|
if torch.isnan(tensor).any():
|
|
print(f"NaN values detected: {name}")
|
|
|
|
def detail(self) -> pd.DataFrame | None:
|
|
"""
|
|
Output MLP model detail information
|
|
|
|
Returns
|
|
-------
|
|
pd.DataFrame
|
|
Feature importance dataframe
|
|
"""
|
|
if not self.fitted:
|
|
logger.info("模型尚未训练，无法显示详细信息")
|
|
return None
|
|
|
|
# 显示模型基本信息
|
|
logger.info(f"输入特征维度: {self.input_size}")
|
|
logger.info(f"隐藏层大小: {self.hidden_sizes}")
|
|
|
|
# 计算模型总参数量
|
|
total_params = sum(p.numel() for p in self.model.parameters())
|
|
logger.info(f"模型总参数量: {total_params:,}")
|
|
|
|
# 显示训练状态信息
|
|
logger.info(f"训练设备: {self.device}")
|
|
logger.info(f"当前学习率: {self.lr}")
|
|
logger.info(f"批次大小: {self.batch_size}")
|
|
|
|
# Calculate feature importance
|
|
importance_df = self._calculate_feature_importance()
|
|
return importance_df
|
|
|
|
def _calculate_feature_importance(self) -> pd.DataFrame:
|
|
"""
|
|
Calculate feature importance
|
|
|
|
Returns
|
|
-------
|
|
pd.DataFrame
|
|
Feature importance dataframe
|
|
"""
|
|
self.model.eval()
|
|
importance_dict: dict[str, float] = {}
|
|
|
|
test_data = torch.randn(1000, self.input_size).to(self.device)
|
|
base_pred = self.model(test_data).detach()
|
|
|
|
noise_level = 0.1
|
|
for i, feature_name in enumerate(self.feature_names):
|
|
perturbed_data = test_data.clone()
|
|
perturbed_data[:, i] += torch.randn(1000).to(self.device) * noise_level
|
|
|
|
with torch.no_grad():
|
|
new_pred = self.model(perturbed_data)
|
|
importance = torch.std(torch.abs(new_pred - base_pred)).item()
|
|
importance_dict[feature_name] = importance
|
|
|
|
df: pd.DataFrame = pd.DataFrame({
|
|
"Feature": list(importance_dict.keys()),
|
|
"Importance": list(importance_dict.values())
|
|
})
|
|
df = df.sort_values("Importance", ascending=False)
|
|
df = df.set_index("Feature")
|
|
|
|
return df
|
|
|
|
|
|
class AverageMeter:
|
|
"""
|
|
Class for calculating and storing average and current values
|
|
|
|
Attributes
|
|
----------
|
|
val : float
|
|
Current value
|
|
avg : float
|
|
Average value
|
|
sum : float
|
|
Sum
|
|
count : int
|
|
Count
|
|
"""
|
|
|
|
def __init__(self) -> None:
|
|
"""
|
|
Initialize AverageMeter
|
|
|
|
Returns
|
|
-------
|
|
None
|
|
"""
|
|
self.reset()
|
|
|
|
def reset(self) -> None:
|
|
"""
|
|
Reset all statistics
|
|
|
|
Returns
|
|
-------
|
|
None
|
|
"""
|
|
self.val: float = 0
|
|
self.avg: float = 0
|
|
self.sum: float = 0
|
|
self.count: int = 0
|
|
|
|
def update(self, val: float, n: int = 1) -> None:
|
|
"""
|
|
Update statistics
|
|
|
|
Parameters
|
|
----------
|
|
val : float
|
|
Current value
|
|
n : int, default 1
|
|
Current batch size
|
|
|
|
Returns
|
|
-------
|
|
None
|
|
"""
|
|
self.val = val
|
|
self.sum += val * n
|
|
self.count += n
|
|
self.avg = self.sum / self.count
|
|
|
|
|
|
class MlpNetwork(nn.Module):
|
|
"""
|
|
Deep Neural Network Model Structure
|
|
|
|
Used to build multi-layer perceptron network structure, supporting multiple hidden layers
|
|
and different activation functions.
|
|
|
|
Attributes
|
|
----------
|
|
network : nn.ModuleList
|
|
List of neural network layers
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
input_size: int,
|
|
output_size: int = 1,
|
|
hidden_sizes: tuple[int] = (256,),
|
|
activation: str = "LeakyReLU"
|
|
) -> None:
|
|
"""
|
|
Constructor
|
|
|
|
Parameters
|
|
----------
|
|
input_size : int
|
|
Input feature dimension, i.e., number of features per sample
|
|
output_size : int, default 1
|
|
Output dimension, used for predicting target values
|
|
hidden_sizes : tuple[int], default (256,)
|
|
Tuple of hidden layer neuron counts, e.g., (256, 128) represents two hidden layers
|
|
with 256 and 128 neurons respectively
|
|
activation : str, default "LeakyReLU"
|
|
Activation function type, options:
|
|
- "LeakyReLU": Leaky ReLU function
|
|
- "SiLU": Sigmoid Linear Unit function
|
|
"""
|
|
super().__init__()
|
|
|
|
# Build network layers
|
|
layers: list[nn.Module] = []
|
|
layer_sizes = [input_size] + list(hidden_sizes)
|
|
|
|
# Input layer Dropout
|
|
layers.append(nn.Dropout(0.05))
|
|
|
|
# Build hidden layers
|
|
for in_size, out_size in zip(layer_sizes[:-1], layer_sizes[1:], strict=False):
|
|
# Add a neural network block: linear layer + batch normalization + activation function
|
|
layers.extend([
|
|
nn.Linear(in_size, out_size),
|
|
nn.BatchNorm1d(out_size),
|
|
self._get_activation(activation)
|
|
])
|
|
|
|
# Output layer
|
|
layers.extend([
|
|
nn.Dropout(0.05),
|
|
nn.Linear(hidden_sizes[-1], output_size)
|
|
])
|
|
|
|
# Combine all layers into a sequence
|
|
self.network = nn.ModuleList(layers)
|
|
|
|
# Initialize network weights
|
|
self._initialize_weights()
|
|
|
|
def _get_activation(self, name: str) -> nn.Module:
|
|
"""
|
|
Get specified activation function layer
|
|
|
|
Parameters
|
|
----------
|
|
name : str
|
|
Activation function name
|
|
|
|
Returns
|
|
-------
|
|
nn.Module
|
|
Activation function layer instance
|
|
|
|
Raises
|
|
------
|
|
ValueError
|
|
When an unsupported activation function type is specified
|
|
"""
|
|
if name == "LeakyReLU":
|
|
return nn.LeakyReLU(negative_slope=0.1)
|
|
elif name == "SiLU":
|
|
return nn.SiLU()
|
|
else:
|
|
raise ValueError(f"Unsupported activation function type: {name}")
|
|
|
|
def _initialize_weights(self) -> None:
|
|
"""
|
|
Initialize network weight parameters
|
|
|
|
Uses Kaiming initialization method for all linear layers, which is particularly
|
|
suitable for deep networks using LeakyReLU activation functions.
|
|
|
|
Returns
|
|
-------
|
|
None
|
|
"""
|
|
for module in self.modules():
|
|
if isinstance(module, nn.Linear):
|
|
nn.init.kaiming_normal_(
|
|
module.weight,
|
|
a=0.1, # LeakyReLU negative slope
|
|
mode="fan_in", # Scale using input node count
|
|
nonlinearity="leaky_relu"
|
|
)
|
|
|
|
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
|
"""
|
|
Forward propagation calculation
|
|
|
|
Parameters
|
|
----------
|
|
x : torch.Tensor
|
|
Input feature tensor, shape (batch_size, input_size)
|
|
|
|
Returns
|
|
-------
|
|
torch.Tensor
|
|
Model output tensor, shape (batch_size, output_size)
|
|
"""
|
|
# Pass through all layers in the network sequentially
|
|
for layer in self.network:
|
|
x = layer(x)
|
|
return x
|