4、数据集构建：从DataFrame到Dataset、训练集/验证集/测试集划分、数据加载器(DataLoader)的使用

好，咱们继续往下走。上一章我们把数据从数据库里捞了出来，整理成了规整的DataFrame。但说实话，DataFrame虽然好用，它毕竟不是为深度学习设计的。你想想看，模型训练的时候要批量取数据、要随机打乱、要并行加载，这些事儿DataFrame干起来就有点力不从心了。

所以这一章，我们就来解决这个问题。我会带你一步步把DataFrame转换成PyTorch的Dataset，再划分出训练集、验证集、测试集，最后用DataLoader把数据喂给模型。嗯，这部分我踩过的坑不少，咱们边聊边避。

4.1 从DataFrame到Dataset：为什么非要转？

我个人习惯，拿到数据后第一件事就是定义自己的Dataset类。为什么？因为Dataset是PyTorch数据管线的起点。它规定了：你的数据长什么样？怎么取一条样本？

在Qlib的选股场景里，一条样本通常包含三样东西：

  • 特征：比如过去N天的动量、波动率、换手率等因子值
  • 标签：未来N天的收益率（或者涨跌分类）
  • 元信息：股票代码、日期等，方便回测时定位

说白了，Dataset就是给模型一个标准接口：你给我一个索引，我返回一个样本。就这么简单。

核心原则：Dataset只负责“如何取一条数据”，不负责“怎么批量取、怎么打乱”。那些是DataLoader的事。

4.2 自定义Dataset：动手写一个

咱们直接上代码。假设你的DataFrame长这样：

import torch
from torch.utils.data import Dataset

class StockDataset(Dataset):
    def __init__(self, df, feature_cols, label_col):
        """
        df: pandas DataFrame，包含特征和标签
        feature_cols: 特征列名列表
        label_col: 标签列名
        """
        self.features = df[feature_cols].values.astype('float32')
        self.labels = df[label_col].values.astype('float32')
        # 保留元信息，方便调试
        self.meta = df[['stock_code', 'date']].values

    def __len__(self):
        return len(self.features)

    def __getitem__(self, idx):
        x = torch.tensor(self.features[idx])
        y = torch.tensor(self.labels[idx]).unsqueeze(0)  # 保持维度一致
        return x, y

这里有个细节我提一下。我在项目中遇到过一个问题：标签的shape不一致。有的模型期望标量，有的期望向量。所以我习惯在__getitem__里统一用unsqueeze(0)，把标量变成[1]的向量。这样后面不管用MSE还是BCE，都不会报shape的错。

小技巧：如果你要做时间序列的滑动窗口，可以在__getitem__里根据idx去切分历史窗口。比如features[idx: idx+window]。但要注意边界处理，别越界。

4.3 训练集/验证集/测试集划分：时间序列的坑

普通机器学习任务，随机划分就完事了。但量化金融不行。为什么？因为时间序列有顺序依赖。你用未来的数据去训练，再用过去的数据去验证，这叫“未来函数”，是量化大忌。

我刚开始做的时候，就犯过这个错。随机划分后模型表现特别好，一上实盘就崩。后来才发现，验证集里混进了未来的信息。嗯，从那以后我再也不敢乱划分了。

正确的做法是按时间切分：

import pandas as pd

# 假设df已经按日期排序
train_end = '2022-12-31'
val_end = '2023-06-30'

train_df = df[df['date'] <= train_end]
val_df = df[(df['date'] > train_end) & (df['date'] <= val_end)]
test_df = df[df['date'] > val_end]

print(f"训练集: {len(train_df)} 条, {train_df['date'].min()} ~ {train_df['date'].max()}")
print(f"验证集: {len(val_df)} 条, {val_df['date'].min()} ~ {val_df['date'].max()}")
print(f"测试集: {len(test_df)} 条, {test_df['date'].min()} ~ {test_df['date'].max()}")

这里有个比例问题。我个人习惯：训练集占60%-70%，验证集和测试集各占15%-20%。但具体要看你的数据时间跨度。如果只有3年数据，那验证集和测试集各给半年就差不多了。

注意：千万不要用随机划分！在量化里，时间顺序就是一切。你想想看，用2023年的数据去预测2022年的收益，这逻辑上就说不通。

4.4 数据加载器(DataLoader)：让数据飞起来

Dataset准备好了，划分也做好了。接下来就是DataLoader登场。它的作用说白了就三件事：批量取数据、打乱顺序、多进程加载。

from torch.utils.data import DataLoader

train_dataset = StockDataset(train_df, feature_cols, label_col)
val_dataset = StockDataset(val_df, feature_cols, label_col)
test_dataset = StockDataset(test_df, feature_cols, label_col)

train_loader = DataLoader(
    train_dataset,
    batch_size=256,      # 每批256条
    shuffle=True,        # 训练集要打乱
    num_workers=4,       # 4个进程并行加载
    pin_memory=True      # 加速GPU传输
)

val_loader = DataLoader(
    val_dataset,
    batch_size=256,
    shuffle=False,       # 验证集不打乱
    num_workers=4
)

test_loader = DataLoader(
    test_dataset,
    batch_size=256,
    shuffle=False
)

这里有几个参数我重点说一下：

  • batch_size：不是越大越好。我试过1024，显存爆了。也试过32，训练太慢。256是个比较稳妥的起点。
  • shuffle：训练集必须打乱，否则模型会学到时间顺序的假规律。验证集和测试集不要打乱，保持顺序方便后续分析。
  • num_workers：一般设成CPU核心数的一半。设太多反而会因为进程切换开销而变慢。
  • pin_memory：如果你用GPU，这个一定要开。它能加速CPU到GPU的数据传输，效果很明显。

避坑指南：我曾经把num_workers设成32，结果数据加载比训练还慢。后来发现是磁盘IO成了瓶颈。如果你的数据在机械硬盘上，num_workers设4-8就够了。SSD可以适当调高。

4.5 核心逻辑流程图

下面这张图，我把整个数据构建的流程画了出来。你可以对照着看，每一步都做了什么。

数据集构建核心流程 原始DataFrame 特征 + 标签 + 元信息 自定义Dataset __len__ + __getitem__ 按时间划分 训练集 (60-70%) 验证集 (15-20%) | 测试集 (15-20%) DataLoader batch_size | shuffle | num_workers → 批量、并行、高效加载 ⚠️ 不要随机划分 时间顺序就是一切 💡 保留元信息 方便回测时定位 ⚡ 开启pin_memory GPU加速必备

4.6 完整代码示例

最后，我把整个流程串起来，给你一个可以直接跑的模板：

import pandas as pd
import torch
from torch.utils.data import Dataset, DataLoader

# 1. 定义Dataset
class StockDataset(Dataset):
    def __init__(self, df, feature_cols, label_col):
        self.features = df[feature_cols].values.astype('float32')
        self.labels = df[label_col].values.astype('float32')
        self.meta = df[['stock_code', 'date']].values

    def __len__(self):
        return len(self.features)

    def __getitem__(self, idx):
        x = torch.tensor(self.features[idx])
        y = torch.tensor(self.labels[idx]).unsqueeze(0)
        return x, y

# 2. 按时间划分
train_df = df[df['date'] <= '2022-12-31']
val_df = df[(df['date'] > '2022-12-31') & (df['date'] <= '2023-06-30')]
test_df = df[df['date'] > '2023-06-30']

# 3. 构建Dataset
feature_cols = ['momentum_5d', 'volatility_10d', 'turnover_20d', 'pe_ratio']
label_col = 'return_5d'

train_dataset = StockDataset(train_df, feature_cols, label_col)
val_dataset = StockDataset(val_df, feature_cols, label_col)
test_dataset = StockDataset(test_df, feature_cols, label_col)

# 4. 构建DataLoader
train_loader = DataLoader(train_dataset, batch_size=256, shuffle=True, num_workers=4, pin_memory=True)
val_loader = DataLoader(val_dataset, batch_size=256, shuffle=False, num_workers=4)
test_loader = DataLoader(test_dataset, batch_size=256, shuffle=False)

# 5. 验证一下
for batch_x, batch_y in train_loader:
    print(f"特征形状: {batch_x.shape}")  # [256, 4]
    print(f"标签形状: {batch_y.shape}")  # [256, 1]
    break

嗯，到这里，数据构建的管线就搭好了。你想想看，从原始的DataFrame，到自定义Dataset，再到按时间划分，最后用DataLoader高效加载——每一步都有它的道理。我在项目中反复调整过这些参数，上面给出的配置是比较通用的起点。你可以根据自己的数据量和硬件情况微调。

下一章，我们会用这个数据管线去训练一个简单的LSTM模型。到时候你就知道，前面这些准备工作有多重要了。

公众号:蓝海资料掘金营，微信deep3321

← 上一章 📖 返回目录 下一章 →