crypto_predictor
1bd9d25f60
- 删除测试脚本: test_2026_classification.py, test_lagllama.py, test_timesfm.py, test_ttm.py - 删除2019相关文件: analyze_2019_*, apply_2019_*, dlinear_2019_*, three_models_2019_* - 添加集成模型相关代码和分析脚本 - 添加N-BEATS模型和CatBoost分类模型
266 lines
8.8 KiB
Python
266 lines
8.8 KiB
Python
"""N-BEATS模型实现：预测未来1小时BTC价格"""
|
||
|
||
import torch
|
||
import torch.nn as nn
|
||
import numpy as np
|
||
|
||
|
||
class NBeatsBlock(nn.Module):
|
||
"""N-BEATS基础块"""
|
||
|
||
def __init__(self, input_size, theta_size, basis_function, layers, layer_size):
|
||
super().__init__()
|
||
self.input_size = input_size
|
||
self.theta_size = theta_size
|
||
self.basis_function = basis_function
|
||
|
||
# 全连接层栈
|
||
self.fc_stack = nn.ModuleList([
|
||
nn.Linear(input_size if i == 0 else layer_size, layer_size)
|
||
for i in range(layers)
|
||
])
|
||
|
||
# Theta层（用于生成基函数参数）
|
||
self.theta_b_fc = nn.Linear(layer_size, theta_size, bias=False) # backcast
|
||
self.theta_f_fc = nn.Linear(layer_size, theta_size, bias=False) # forecast
|
||
|
||
self.relu = nn.ReLU()
|
||
|
||
def forward(self, x):
|
||
"""
|
||
Args:
|
||
x: (batch_size, input_size)
|
||
Returns:
|
||
backcast: (batch_size, input_size)
|
||
forecast: (batch_size, forecast_size)
|
||
"""
|
||
# 通过FC层栈
|
||
for fc in self.fc_stack:
|
||
x = self.relu(fc(x))
|
||
|
||
# 生成theta参数
|
||
theta_b = self.theta_b_fc(x) # backcast coefficients
|
||
theta_f = self.theta_f_fc(x) # forecast coefficients
|
||
|
||
# 应用基函数
|
||
backcast = self.basis_function(theta_b, is_forecast=False)
|
||
forecast = self.basis_function(theta_f, is_forecast=True)
|
||
|
||
return backcast, forecast
|
||
|
||
|
||
class GenericBasis(nn.Module):
|
||
"""通用基函数"""
|
||
|
||
def __init__(self, backcast_size, forecast_size):
|
||
super().__init__()
|
||
self.backcast_size = backcast_size
|
||
self.forecast_size = forecast_size
|
||
|
||
def forward(self, theta, is_forecast):
|
||
if is_forecast:
|
||
return theta[:, :self.forecast_size]
|
||
else:
|
||
return theta[:, :self.backcast_size]
|
||
|
||
|
||
class TrendBasis(nn.Module):
|
||
"""趋势基函数（多项式）"""
|
||
|
||
def __init__(self, degree, backcast_size, forecast_size):
|
||
super().__init__()
|
||
self.backcast_size = backcast_size
|
||
self.forecast_size = forecast_size
|
||
|
||
# 多项式时间点
|
||
backcast_time = torch.arange(backcast_size, dtype=torch.float32) / backcast_size
|
||
forecast_time = torch.arange(forecast_size, dtype=torch.float32) / forecast_size
|
||
|
||
# 构建多项式基
|
||
self.backcast_basis = torch.stack([backcast_time ** i for i in range(degree)], dim=1)
|
||
self.forecast_basis = torch.stack([forecast_time ** i for i in range(degree)], dim=1)
|
||
|
||
def forward(self, theta, is_forecast):
|
||
basis = self.forecast_basis if is_forecast else self.backcast_basis
|
||
basis = basis.to(theta.device)
|
||
return torch.matmul(theta, basis.T)
|
||
|
||
|
||
class SeasonalityBasis(nn.Module):
|
||
"""季节性基函数（傅里叶级数）"""
|
||
|
||
def __init__(self, harmonics, backcast_size, forecast_size):
|
||
super().__init__()
|
||
self.backcast_size = backcast_size
|
||
self.forecast_size = forecast_size
|
||
|
||
backcast_time = 2 * np.pi * torch.arange(backcast_size, dtype=torch.float32) / backcast_size
|
||
forecast_time = 2 * np.pi * torch.arange(forecast_size, dtype=torch.float32) / forecast_size
|
||
|
||
# 构建傅里叶基（sin和cos）
|
||
backcast_basis = []
|
||
forecast_basis = []
|
||
for i in range(1, harmonics + 1):
|
||
backcast_basis.extend([torch.cos(i * backcast_time), torch.sin(i * backcast_time)])
|
||
forecast_basis.extend([torch.cos(i * forecast_time), torch.sin(i * forecast_time)])
|
||
|
||
self.backcast_basis = torch.stack(backcast_basis, dim=1)
|
||
self.forecast_basis = torch.stack(forecast_basis, dim=1)
|
||
|
||
def forward(self, theta, is_forecast):
|
||
basis = self.forecast_basis if is_forecast else self.backcast_basis
|
||
basis = basis.to(theta.device)
|
||
return torch.matmul(theta, basis.T)
|
||
|
||
|
||
class NBeatsStack(nn.Module):
|
||
"""N-BEATS栈（多个块的组合）"""
|
||
|
||
def __init__(self, input_size, forecast_size, num_blocks, block_layers,
|
||
layer_size, theta_size, basis_function):
|
||
super().__init__()
|
||
self.blocks = nn.ModuleList([
|
||
NBeatsBlock(
|
||
input_size=input_size,
|
||
theta_size=theta_size,
|
||
basis_function=basis_function,
|
||
layers=block_layers,
|
||
layer_size=layer_size
|
||
)
|
||
for _ in range(num_blocks)
|
||
])
|
||
self.forecast_size = forecast_size
|
||
|
||
def forward(self, x):
|
||
"""
|
||
Args:
|
||
x: (batch_size, input_size)
|
||
Returns:
|
||
forecast: (batch_size, forecast_size)
|
||
"""
|
||
residual = x
|
||
forecast = torch.zeros(x.size(0), self.forecast_size, device=x.device)
|
||
|
||
for block in self.blocks:
|
||
backcast, block_forecast = block(residual)
|
||
residual = residual - backcast
|
||
forecast = forecast + block_forecast
|
||
|
||
return forecast
|
||
|
||
|
||
class NBeatsConfig:
|
||
"""N-BEATS模型配置"""
|
||
|
||
def __init__(
|
||
self,
|
||
input_size=168, # 输入序列长度
|
||
forecast_size=1, # 预测长度
|
||
stacks=['trend', 'seasonality', 'generic'], # 使用的栈类型
|
||
num_blocks_per_stack=3, # 每个栈的块数
|
||
block_layers=4, # 每个块的FC层数
|
||
layer_size=256, # FC层大小
|
||
trend_degree=2, # 趋势多项式阶数
|
||
seasonality_harmonics=4, # 季节性谐波数
|
||
):
|
||
self.input_size = input_size
|
||
self.forecast_size = forecast_size
|
||
self.stacks = stacks
|
||
self.num_blocks_per_stack = num_blocks_per_stack
|
||
self.block_layers = block_layers
|
||
self.layer_size = layer_size
|
||
self.trend_degree = trend_degree
|
||
self.seasonality_harmonics = seasonality_harmonics
|
||
|
||
|
||
class NBeats(nn.Module):
|
||
"""N-BEATS模型"""
|
||
|
||
def __init__(self, config):
|
||
super().__init__()
|
||
self.config = config
|
||
|
||
# 创建多个栈
|
||
self.stacks = nn.ModuleList()
|
||
|
||
for stack_type in config.stacks:
|
||
if stack_type == 'trend':
|
||
basis_function = TrendBasis(
|
||
degree=config.trend_degree,
|
||
backcast_size=config.input_size,
|
||
forecast_size=config.forecast_size
|
||
)
|
||
theta_size = config.trend_degree
|
||
|
||
elif stack_type == 'seasonality':
|
||
basis_function = SeasonalityBasis(
|
||
harmonics=config.seasonality_harmonics,
|
||
backcast_size=config.input_size,
|
||
forecast_size=config.forecast_size
|
||
)
|
||
theta_size = 2 * config.seasonality_harmonics
|
||
|
||
else: # generic
|
||
basis_function = GenericBasis(
|
||
backcast_size=config.input_size,
|
||
forecast_size=config.forecast_size
|
||
)
|
||
theta_size = max(config.input_size, config.forecast_size)
|
||
|
||
stack = NBeatsStack(
|
||
input_size=config.input_size,
|
||
forecast_size=config.forecast_size,
|
||
num_blocks=config.num_blocks_per_stack,
|
||
block_layers=config.block_layers,
|
||
layer_size=config.layer_size,
|
||
theta_size=theta_size,
|
||
basis_function=basis_function
|
||
)
|
||
self.stacks.append(stack)
|
||
|
||
def forward(self, x):
|
||
"""
|
||
Args:
|
||
x: (batch_size, lookback, features) 或 (batch_size, lookback)
|
||
Returns:
|
||
forecast: (batch_size, forecast_size)
|
||
"""
|
||
# 如果输入是多特征，取最后一个特征（close price）
|
||
if len(x.shape) == 3:
|
||
x = x[:, :, -1] # (batch_size, lookback)
|
||
|
||
# 确保是2D
|
||
if len(x.shape) == 1:
|
||
x = x.unsqueeze(0)
|
||
|
||
# 通过所有栈
|
||
forecast = torch.zeros(x.size(0), self.config.forecast_size, device=x.device)
|
||
residual = x
|
||
|
||
for stack in self.stacks:
|
||
stack_forecast = stack(residual)
|
||
forecast = forecast + stack_forecast
|
||
# 更新残差（用于下一个栈）
|
||
# 注意：这里简化了，标准N-BEATS会从每个栈获取backcast
|
||
|
||
return forecast.squeeze(-1) # (batch_size,)
|
||
|
||
def count_parameters(self):
|
||
"""统计模型参数数量"""
|
||
return sum(p.numel() for p in self.parameters() if p.requires_grad)
|
||
|
||
def save(self, path):
|
||
"""保存模型"""
|
||
torch.save({
|
||
'model_state_dict': self.state_dict(),
|
||
'config': self.config
|
||
}, path)
|
||
|
||
@classmethod
|
||
def load(cls, path, device='cpu'):
|
||
"""加载模型"""
|
||
checkpoint = torch.load(path, map_location=device)
|
||
model = cls(checkpoint['config'])
|
||
model.load_state_dict(checkpoint['model_state_dict'])
|
||
return model
|