liaozhaorun
36a3ccbcc8
- 添加 TabMModel、TabPFNModel 深度学习模型实现 - 新增 DataQualityAnalyzer 进行训练前数据质量诊断 - 改进数据处理器 NaN/null 双重处理，增强数据鲁棒性 - 支持 train_skip_days 参数跳过训练初期数据不足期 - Pipeline 自动清理标签为 NaN 的样本
520 lines
19 KiB
Python
520 lines
19 KiB
Python
"""TabM NaN 问题诊断测试 + train_skip_days 功能验证
|
||
|
||
诊断 loss 为 nan 的根因：
|
||
1. 标签中是否有 NaN 或极端值
|
||
2. 标准化后是否有 NaN
|
||
3. 是否有方差为0或接近0的列
|
||
4. train_skip_days 功能是否正常工作
|
||
"""
|
||
|
||
import numpy as np
|
||
import polars as pl
|
||
import pytest
|
||
|
||
from src.factors import FactorEngine
|
||
from src.training import (
|
||
FactorManager,
|
||
DataPipeline,
|
||
TabMRegressionTask,
|
||
NullFiller,
|
||
Winsorizer,
|
||
StandardScaler,
|
||
)
|
||
from src.training.components.filters import STFilter
|
||
from src.experiment.common import (
|
||
SELECTED_FACTORS,
|
||
FACTOR_DEFINITIONS,
|
||
LABEL_NAME,
|
||
LABEL_FACTOR,
|
||
stock_pool_filter,
|
||
STOCK_FILTER_REQUIRED_COLUMNS,
|
||
)
|
||
|
||
|
||
# TabM 模型参数（简化用于快速测试）
|
||
MODEL_PARAMS = {
|
||
"n_blocks": 2,
|
||
"d_block": 128,
|
||
"dropout": 0.1,
|
||
"ensemble_size": 4, # 简化
|
||
"batch_size": 256,
|
||
"learning_rate": 1e-4, # 降低学习率
|
||
"weight_decay": 1e-5,
|
||
"epochs": 3,
|
||
"early_stopping_patience": 5,
|
||
}
|
||
|
||
# 测试日期范围
|
||
TEST_DATE_RANGE = {
|
||
"train": ("20200101", "20200630"),
|
||
"val": ("20200701", "20200731"),
|
||
"test": ("20200801", "20200831"),
|
||
}
|
||
|
||
# 小范围测试用的因子排除列表
|
||
EXCLUDED_FACTORS = [
|
||
"GTJA_alpha001",
|
||
"GTJA_alpha002",
|
||
"GTJA_alpha003",
|
||
"GTJA_alpha004",
|
||
"GTJA_alpha005",
|
||
"GTJA_alpha006",
|
||
"GTJA_alpha007",
|
||
"GTJA_alpha008",
|
||
"GTJA_alpha009",
|
||
"GTJA_alpha010",
|
||
]
|
||
|
||
|
||
class TestTabMNanDebug:
|
||
"""TabM NaN 问题诊断测试类（使用 DataPipeline）"""
|
||
|
||
@pytest.fixture(scope="class")
|
||
def engine_and_factor_manager(self):
|
||
"""准备 FactorEngine 和 FactorManager"""
|
||
engine = FactorEngine()
|
||
|
||
factor_manager = FactorManager(
|
||
selected_factors=SELECTED_FACTORS,
|
||
factor_definitions=FACTOR_DEFINITIONS,
|
||
label_factor=LABEL_FACTOR,
|
||
excluded_factors=EXCLUDED_FACTORS,
|
||
)
|
||
|
||
return {
|
||
"engine": engine,
|
||
"factor_manager": factor_manager,
|
||
}
|
||
|
||
@pytest.fixture(scope="class")
|
||
def pipeline_data(self, engine_and_factor_manager):
|
||
"""使用 DataPipeline 准备数据，并过滤掉全为NaN的列"""
|
||
engine = engine_and_factor_manager["engine"]
|
||
factor_manager = engine_and_factor_manager["factor_manager"]
|
||
|
||
# 创建 DataPipeline（使用 train_skip_days=0 进行基础测试）
|
||
pipeline = DataPipeline(
|
||
factor_manager=factor_manager,
|
||
processor_configs=[
|
||
(NullFiller, {"strategy": "mean"}),
|
||
(Winsorizer, {"lower": 0.01, "upper": 0.99}),
|
||
(StandardScaler, {}),
|
||
],
|
||
filters=[STFilter(data_router=engine.router)],
|
||
stock_pool_filter_func=stock_pool_filter,
|
||
stock_pool_required_columns=STOCK_FILTER_REQUIRED_COLUMNS,
|
||
train_skip_days=0, # 先不跳过天数进行测试
|
||
)
|
||
|
||
# 准备数据
|
||
data = pipeline.prepare_data(
|
||
engine=engine,
|
||
date_range=TEST_DATE_RANGE,
|
||
label_name=LABEL_NAME,
|
||
verbose=True,
|
||
)
|
||
|
||
# 获取特征列
|
||
feature_cols = factor_manager.register_to_engine(engine, verbose=False)
|
||
|
||
# 过滤掉全为NaN的列（这些因子计算失败，没有修复意义）
|
||
print("\n[DataPipeline] 检查并过滤全为NaN的特征列...")
|
||
for split_name in ["train", "val", "test"]:
|
||
X_df = data[split_name]["X"]
|
||
y_series = data[split_name]["y"]
|
||
raw_data = data[split_name]["raw_data"]
|
||
|
||
# 删除标签为NaN的行
|
||
y_nan_count = y_series.null_count()
|
||
if y_nan_count > 0:
|
||
print(f" {split_name}: 发现 {y_nan_count} 个标签为NaN的行，将被删除")
|
||
# 创建有效标签的mask
|
||
valid_mask = y_series.is_not_null()
|
||
# 过滤所有相关数据
|
||
X_df = X_df.filter(valid_mask)
|
||
y_series = y_series.filter(valid_mask)
|
||
raw_data = raw_data.filter(valid_mask)
|
||
# 更新数据
|
||
data[split_name]["X"] = X_df
|
||
data[split_name]["y"] = y_series
|
||
data[split_name]["raw_data"] = raw_data
|
||
|
||
# 检查每列的NaN数量
|
||
nan_counts = {col: X_df[col].null_count() for col in X_df.columns}
|
||
total_rows = len(X_df)
|
||
|
||
# 找出全为NaN的列
|
||
all_nan_cols = [
|
||
col for col, count in nan_counts.items() if count == total_rows
|
||
]
|
||
|
||
if all_nan_cols:
|
||
print(
|
||
f" {split_name}: 发现 {len(all_nan_cols)} 个全为NaN的列，将被删除"
|
||
)
|
||
print(
|
||
f" 列名: {all_nan_cols[:5]}{'...' if len(all_nan_cols) > 5 else ''}"
|
||
)
|
||
|
||
# 更新feature_cols
|
||
feature_cols = [c for c in feature_cols if c not in all_nan_cols]
|
||
|
||
# 从X中删除这些列
|
||
X_df = X_df.select(feature_cols)
|
||
|
||
# 从raw_data中也删除这些列（保留原始特征列以外的列如trade_date, ts_code等）
|
||
raw_cols_to_keep = [
|
||
c for c in raw_data.columns if c not in all_nan_cols
|
||
]
|
||
raw_data = raw_data.select(raw_cols_to_keep)
|
||
|
||
# 更新数据
|
||
data[split_name]["X"] = X_df
|
||
data[split_name]["raw_data"] = raw_data
|
||
data[split_name]["feature_cols"] = feature_cols
|
||
|
||
# 验证没有NaN了
|
||
X_np = X_df.to_numpy()
|
||
nan_count = np.isnan(X_np).sum()
|
||
assert nan_count == 0, f"{split_name} 中仍有 {nan_count} 个NaN"
|
||
|
||
print(f" 过滤后特征数: {len(feature_cols)}")
|
||
print(" [通过] 所有特征列均无NaN")
|
||
|
||
return {
|
||
"pipeline": pipeline,
|
||
"data": data,
|
||
"feature_cols": feature_cols,
|
||
"engine": engine,
|
||
}
|
||
|
||
def test_pipeline_data_structure(self, pipeline_data):
|
||
"""诊断0: DataPipeline 返回的数据结构检查"""
|
||
data = pipeline_data["data"]
|
||
feature_cols = pipeline_data["feature_cols"]
|
||
|
||
print("\n[诊断0] DataPipeline 数据结构检查:")
|
||
|
||
# 检查数据结构
|
||
assert "train" in data, "缺少 train 数据"
|
||
assert "val" in data, "缺少 val 数据"
|
||
assert "test" in data, "缺少 test 数据"
|
||
|
||
for split_name in ["train", "val", "test"]:
|
||
split_data = data[split_name]
|
||
assert "X" in split_data, f"{split_name} 缺少 X"
|
||
assert "y" in split_data, f"{split_name} 缺少 y"
|
||
assert "raw_data" in split_data, f"{split_name} 缺少 raw_data"
|
||
assert "feature_cols" in split_data, f"{split_name} 缺少 feature_cols"
|
||
|
||
X = split_data["X"]
|
||
y = split_data["y"]
|
||
|
||
print(f" {split_name}:")
|
||
print(f" X 形状: {X.shape}")
|
||
print(f" y 形状: {len(y)}")
|
||
print(f" 特征数: {len(split_data['feature_cols'])}")
|
||
|
||
# 检查维度一致性
|
||
assert X.shape[0] == len(y), f"{split_name} X 和 y 行数不匹配"
|
||
assert X.shape[1] == len(feature_cols), f"{split_name} 特征数不匹配"
|
||
|
||
print(" [通过] 数据结构正确")
|
||
|
||
def test_label_quality_with_pipeline(self, pipeline_data):
|
||
"""诊断1: 使用 DataPipeline 后的标签数据质量检查"""
|
||
data = pipeline_data["data"]
|
||
|
||
# 检查所有数据集的标签质量
|
||
for split_name in ["train", "val", "test"]:
|
||
y = data[split_name]["y"]
|
||
y_np = y.to_numpy()
|
||
|
||
print(f"\n[诊断1-{split_name}] 标签数据质量:")
|
||
print(f" 总数: {len(y)}")
|
||
print(f" NaN数量: {y.null_count()}")
|
||
print(f" 均值: {y.mean():.6f}")
|
||
print(f" 标准差: {y.std():.6f}")
|
||
print(f" 最小值: {y.min():.6f}")
|
||
print(f" 最大值: {y.max():.6f}")
|
||
|
||
inf_count = np.isinf(y_np).sum()
|
||
nan_count = np.isnan(y_np).sum()
|
||
print(f" inf数量: {inf_count}")
|
||
print(f" nan数量: {nan_count}")
|
||
|
||
assert inf_count == 0, f"{split_name} 标签含 inf: {inf_count}"
|
||
assert nan_count == 0, f"{split_name} 标签含 nan: {nan_count}"
|
||
|
||
def test_processed_data_quality_with_pipeline(self, pipeline_data):
|
||
"""诊断2: 使用 DataPipeline 处理后的特征数据质量"""
|
||
data = pipeline_data["data"]
|
||
feature_cols = pipeline_data["feature_cols"]
|
||
|
||
for split_name in ["train", "val", "test"]:
|
||
X = data[split_name]["X"]
|
||
y = data[split_name]["y"]
|
||
|
||
X_np = X.to_numpy().astype(np.float32)
|
||
y_np = y.to_numpy().astype(np.float32)
|
||
|
||
print(f"\n[诊断2-{split_name}] 处理后数据质量:")
|
||
print(f" X 形状: {X_np.shape}, dtype: {X_np.dtype}")
|
||
print(f" y 形状: {y_np.shape}, dtype: {y_np.dtype}")
|
||
print(f" X中NaN: {np.isnan(X_np).sum()}")
|
||
print(f" X中Inf: {np.isinf(X_np).sum()}")
|
||
print(f" y中NaN: {np.isnan(y_np).sum()}")
|
||
print(f" y中Inf: {np.isinf(y_np).sum()}")
|
||
|
||
assert np.isnan(X_np).sum() == 0, f"{split_name} X含NaN"
|
||
assert np.isnan(y_np).sum() == 0, f"{split_name} y含NaN"
|
||
assert np.isinf(X_np).sum() == 0, f"{split_name} X含Inf"
|
||
assert np.isinf(y_np).sum() == 0, f"{split_name} y含Inf"
|
||
|
||
def test_train_skip_days_functionality(self, engine_and_factor_manager):
|
||
"""诊断3: train_skip_days 功能验证"""
|
||
engine = engine_and_factor_manager["engine"]
|
||
factor_manager = engine_and_factor_manager["factor_manager"]
|
||
|
||
print("\n[诊断3] train_skip_days 功能验证:")
|
||
|
||
# 创建一个跳过50天数据的 pipeline
|
||
skip_days = 50
|
||
pipeline_with_skip = DataPipeline(
|
||
factor_manager=factor_manager,
|
||
processor_configs=[
|
||
(NullFiller, {"strategy": "mean"}),
|
||
(Winsorizer, {"lower": 0.01, "upper": 0.99}),
|
||
(StandardScaler, {}),
|
||
],
|
||
filters=[STFilter(data_router=engine.router)],
|
||
stock_pool_filter_func=stock_pool_filter,
|
||
stock_pool_required_columns=STOCK_FILTER_REQUIRED_COLUMNS,
|
||
train_skip_days=skip_days,
|
||
)
|
||
|
||
# 准备数据
|
||
data_with_skip = pipeline_with_skip.prepare_data(
|
||
engine=engine,
|
||
date_range=TEST_DATE_RANGE,
|
||
label_name=LABEL_NAME,
|
||
verbose=True,
|
||
)
|
||
|
||
# 获取原始数据（不跳过天数）用于对比
|
||
pipeline_no_skip = DataPipeline(
|
||
factor_manager=factor_manager,
|
||
processor_configs=[
|
||
(NullFiller, {"strategy": "mean"}),
|
||
(Winsorizer, {"lower": 0.01, "upper": 0.99}),
|
||
(StandardScaler, {}),
|
||
],
|
||
filters=[STFilter(data_router=engine.router)],
|
||
stock_pool_filter_func=stock_pool_filter,
|
||
stock_pool_required_columns=STOCK_FILTER_REQUIRED_COLUMNS,
|
||
train_skip_days=0,
|
||
)
|
||
|
||
data_no_skip = pipeline_no_skip.prepare_data(
|
||
engine=engine,
|
||
date_range=TEST_DATE_RANGE,
|
||
label_name=LABEL_NAME,
|
||
verbose=False,
|
||
)
|
||
|
||
# 验证训练数据减少了
|
||
train_with_skip = data_with_skip["train"]["raw_data"]
|
||
train_no_skip = data_no_skip["train"]["raw_data"]
|
||
|
||
print(f"\n 对比结果:")
|
||
print(f" 不跳过时训练数据: {len(train_no_skip)} 条")
|
||
print(f" 跳过{skip_days}天后: {len(train_with_skip)} 条")
|
||
print(f" 减少: {len(train_no_skip) - len(train_with_skip)} 条")
|
||
|
||
# 验证验证集和测试集不受影响
|
||
assert len(data_with_skip["val"]["raw_data"]) == len(
|
||
data_no_skip["val"]["raw_data"]
|
||
), "val 数据不应该受 train_skip_days 影响"
|
||
assert len(data_with_skip["test"]["raw_data"]) == len(
|
||
data_no_skip["test"]["raw_data"]
|
||
), "test 数据不应该受 train_skip_days 影响"
|
||
|
||
# 验证日期确实被跳过了
|
||
if len(train_no_skip) > 0:
|
||
dates_no_skip = sorted(train_no_skip["trade_date"].unique())
|
||
dates_with_skip = sorted(train_with_skip["trade_date"].unique())
|
||
|
||
print(f"\n 日期对比:")
|
||
print(
|
||
f" 不跳过 - 最早日期: {dates_no_skip[0]}, 共 {len(dates_no_skip)} 个交易日"
|
||
)
|
||
print(
|
||
f" 跳过 - 最早日期: {dates_with_skip[0]}, 共 {len(dates_with_skip)} 个交易日"
|
||
)
|
||
|
||
# 验证跳过的数据确实从更晚的日期开始
|
||
if len(dates_no_skip) > skip_days:
|
||
expected_start_date = dates_no_skip[skip_days]
|
||
assert dates_with_skip[0] == expected_start_date, (
|
||
f"预期从 {expected_start_date} 开始，实际从 {dates_with_skip[0]} 开始"
|
||
)
|
||
print(f" [通过] 正确跳过前 {skip_days} 个交易日")
|
||
|
||
def test_training_with_pipeline(self, pipeline_data):
|
||
"""诊断4: 使用 DataPipeline 处理后的数据进行训练测试"""
|
||
import torch
|
||
from torch.utils.data import DataLoader, TensorDataset
|
||
from tabm import TabM
|
||
|
||
data = pipeline_data["data"]
|
||
feature_cols = pipeline_data["feature_cols"]
|
||
|
||
# 获取训练数据
|
||
X_train_df = data["train"]["X"]
|
||
y_train_series = data["train"]["y"]
|
||
|
||
# 删除标签为NaN的行
|
||
valid_mask = y_train_series.is_not_null()
|
||
y_nan_count = y_train_series.null_count()
|
||
if y_nan_count > 0:
|
||
print(f" 发现 {y_nan_count} 个标签为NaN的行，将被删除")
|
||
X_train_df = X_train_df.filter(valid_mask)
|
||
y_train_series = y_train_series.filter(valid_mask)
|
||
|
||
X_train = X_train_df.to_numpy().astype(np.float32)
|
||
y_train = y_train_series.to_numpy().astype(np.float32)
|
||
|
||
# 只取前1000条加速测试
|
||
X_train = X_train[:1000]
|
||
y_train = y_train[:1000]
|
||
|
||
print(f"\n[诊断4] DataPipeline 数据训练测试:")
|
||
print(f" X 形状: {X_train.shape}")
|
||
print(f" y 形状: {y_train.shape}")
|
||
print(f" X中NaN: {np.isnan(X_train).sum()}, Inf: {np.isinf(X_train).sum()}")
|
||
print(f" y中NaN: {np.isnan(y_train).sum()}, Inf: {np.isinf(y_train).sum()}")
|
||
|
||
# 创建 TabM 模型
|
||
n_features = X_train.shape[1]
|
||
model = TabM.make(
|
||
n_num_features=n_features,
|
||
cat_cardinalities=[],
|
||
d_out=1,
|
||
n_blocks=2,
|
||
d_block=128,
|
||
dropout=0.1,
|
||
k=4,
|
||
)
|
||
|
||
# 训练一个 epoch
|
||
dataset = TensorDataset(torch.from_numpy(X_train), torch.from_numpy(y_train))
|
||
loader = DataLoader(dataset, batch_size=256, shuffle=True)
|
||
|
||
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)
|
||
criterion = torch.nn.MSELoss()
|
||
|
||
model.train()
|
||
losses = []
|
||
|
||
for batch_idx, (bx, by) in enumerate(loader):
|
||
optimizer.zero_grad()
|
||
outputs = model(bx) # [B, E, 1]
|
||
outputs_squeezed = outputs.squeeze(-1) # [B, E]
|
||
by_expanded = by.unsqueeze(-1).expand(-1, 4) # [B, E]
|
||
loss = criterion(outputs_squeezed, by_expanded)
|
||
loss.backward()
|
||
optimizer.step()
|
||
|
||
losses.append(loss.item())
|
||
print(f" 批次 {batch_idx + 1} loss: {loss.item():.6f}")
|
||
|
||
assert not torch.isnan(loss), f"批次 {batch_idx + 1} loss 为 nan!"
|
||
|
||
if batch_idx >= 2: # 只测试前3个批次
|
||
break
|
||
|
||
print(f" [通过] 所有批次 loss 正常，无 NaN")
|
||
|
||
|
||
class TestTrainSkipDaysEdgeCases:
|
||
"""train_skip_days 边界情况测试"""
|
||
|
||
@pytest.fixture(scope="class")
|
||
def engine_and_factor_manager(self):
|
||
"""准备 FactorEngine 和 FactorManager"""
|
||
engine = FactorEngine()
|
||
|
||
factor_manager = FactorManager(
|
||
selected_factors=SELECTED_FACTORS,
|
||
factor_definitions=FACTOR_DEFINITIONS,
|
||
label_factor=LABEL_FACTOR,
|
||
excluded_factors=EXCLUDED_FACTORS,
|
||
)
|
||
|
||
return {
|
||
"engine": engine,
|
||
"factor_manager": factor_manager,
|
||
}
|
||
|
||
def test_skip_more_than_available_days(self, engine_and_factor_manager):
|
||
"""测试跳过天数超过可用天数的情况"""
|
||
engine = engine_and_factor_manager["engine"]
|
||
factor_manager = engine_and_factor_manager["factor_manager"]
|
||
|
||
print("\n[边界测试] 跳过天数超过可用天数:")
|
||
|
||
# 使用一个很大的跳过天数
|
||
pipeline = DataPipeline(
|
||
factor_manager=factor_manager,
|
||
processor_configs=[(NullFiller, {"strategy": "mean"})],
|
||
filters=[STFilter(data_router=engine.router)],
|
||
stock_pool_filter_func=stock_pool_filter,
|
||
stock_pool_required_columns=STOCK_FILTER_REQUIRED_COLUMNS,
|
||
train_skip_days=1000, # 超过测试期间的交易日数
|
||
)
|
||
|
||
# 准备数据（应该能正常运行并发出警告）
|
||
data = pipeline.prepare_data(
|
||
engine=engine,
|
||
date_range=TEST_DATE_RANGE,
|
||
label_name=LABEL_NAME,
|
||
verbose=True,
|
||
)
|
||
|
||
# 即使跳过天数很多，也应该有数据返回
|
||
# 如果交易日数少于跳过天数，应该保留所有数据（只发出警告）
|
||
train_data = data["train"]["raw_data"]
|
||
print(f" 训练数据量: {len(train_data)} 条")
|
||
print(f" [通过] 程序未崩溃")
|
||
|
||
def test_skip_zero_days(self, engine_and_factor_manager):
|
||
"""测试跳过0天（即不跳过）"""
|
||
engine = engine_and_factor_manager["engine"]
|
||
factor_manager = engine_and_factor_manager["factor_manager"]
|
||
|
||
print("\n[边界测试] 跳过0天:")
|
||
|
||
pipeline = DataPipeline(
|
||
factor_manager=factor_manager,
|
||
processor_configs=[(NullFiller, {"strategy": "mean"})],
|
||
filters=[STFilter(data_router=engine.router)],
|
||
stock_pool_filter_func=stock_pool_filter,
|
||
stock_pool_required_columns=STOCK_FILTER_REQUIRED_COLUMNS,
|
||
train_skip_days=0,
|
||
)
|
||
|
||
data = pipeline.prepare_data(
|
||
engine=engine,
|
||
date_range=TEST_DATE_RANGE,
|
||
label_name=LABEL_NAME,
|
||
verbose=False,
|
||
)
|
||
|
||
train_data = data["train"]["raw_data"]
|
||
print(f" 训练数据量: {len(train_data)} 条")
|
||
print(f" [通过] skip=0 时数据正常")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
pytest.main([__file__, "-v", "-s"])
|