AI DataLoader

DolphinDB Python API 自 1.30.22.2 版本起提供深度学习工具类 DDBDataLoader，提供对 DolphinDB SQL 对应的数据集进行批量拆分和重新洗牌的易用接口，将 DolphinDB 中的数据直接对接到 PyTorch 中。

pip install dolphindb-tools 

期待输出：

注意：

1. 若 sql 查询的结果表中包含不支持的类型，即便其列名被包含在 targetCol 中，即表示迭代中 y 对应的列名，详细见接口说明，该数据列也不会出现在输入数据和目标数据中。
2. 支持上述类型的 ArrayVector 类型。如果使用 ArrayVector 列，需要保证输入数据或目标数据全部为 ArrayVector 类型。
3. torch.bool 不支持布尔型数据的空值，因此获取 BOOL 类型数据前需确保不包含空值。

提供 DDBDataLoader 类来加载和访问数据，接口如下：

DDBDataLoader(
    ddbSession: Session,
    sql: str,
    targetCol: List[str],
    batchSize: int = 1,
    shuffle: bool = True,
    windowSize: Union[List[int], int, None] = None,
    windowStride: Union[List[int], int, None] = None,
    *,
    inputCol: Optional[List[str]] = None,
    excludeCol: Optional[List[str]] = None,
    repartitionCol: str = None,
    repartitionScheme: List[str] = None,
    groupCol: str = None,   
    groupScheme: List[str] = None,
    seed: Optional[int] = None,
    dropLast: bool = False,
    offset: int = None,
    device: str = "cpu",
    prefetchBatch: int = 1,
    prepartitionNum: int = 2,
    groupPoolSize: int = 3,
    **kwargs
)

服务端

• 硬件环境

  硬件名称	配置信息
  主机名	HostName
  外网 IP	xxx.xxx.xxx.218
  操作系统	Linux（内核版本3.10以上）
  内存	500 GB
  CPU	x86_64（64核心）
  GPU	A100
  网络	万兆以太网

• 软件环境

  软件名称	版本信息
  DolphinDB	2.00.10.1
  ddbtools	0.1a1
  python	3.8.17
  dolphindb	1.30.22.2
  numpy，torch ，pandas	1.24.4, 2.0.0+cu118, 1.5.2

• 性能测试工具

  使用 Python 中的第三方库 line_profiler （4.0.3），将待测试代码封装为函数后添加 @profile 装饰器，在终端执行 kernprof -l -v test.py 进行性能测试。

• 测试数据

  快照 3 秒频因子数据，生成总数据约为 277G，测试数据生成脚本如下：

  在 DolphinDB 客户端执行，指定 Datetime 和 Symbol 为分区列和排序列，在数据库 dfs://test_ai_dataloader 中创建分区表 wide_factor_table。表中包含 Datetime 时间列和 Symbol 股票名称列，以及 1000 列因子列（名称从 f000001 到 f001000）。类型分别为 DATETIME 和 SYMBOL，因子列类型全部使用 DOUBLE。详细代码见工程代码中 ddb_scripts.dos，核心代码如下：

  dbName = "dfs://test_ai_dataloader"
  tbName = "wide_factor_table"
  
  if (existsDatabase(dbName)) {
      dropDatabase(dbName)
  }
  
  // 股票数量
  numSymbols = 250
  // 因子数量
  numFactors = 1000
  
  dateBegin = 2020.01.01
  dateEnd = 2020.01.31
  symbolList = symbol(lpad(string(1..numSymbols), 6, "0") + ".SH")
  factorList = lpad(string(1..numFactors), 7,"f000000")
  
  colNames = ["Datetime", "Symbol"] join factorList
  colTypes = [DATETIME, SYMBOL] join take(DOUBLE, numFactors)
  schema = table(1:0, colNames, colTypes)

  写入完成后，使用下面的脚本打印 SQL 查询结果，确认已经写入成功。

  select DateTime, Symbol, f000001 from loadTable("dfs://test_ai_dataloader", "wide_factor_table") where Symbol=`000001.SH, date(DateTime)=2020.01.31

  示例数据如下：

  
  
  

在传统深度学习中，通常会采取以下步骤来处理训练数据：

• 生成二进制数据文件
• 生成索引信息 pkl 文件
• 使用 PyTorch DataLoader 方式加载数据到模型中

首先使用 numpy 库生成二进制数据文件，此阶段耗时约为 83分钟，详细见 prepare_data.py，核心代码如下:

st = time.time()
for symbol in symbols:
    for t in times:
        sql_tmp = sql + f""" where Symbol={symbol}, date(DateTime)={t}"""
        data = s.run(sql_tmp, pickleTableToList=True)
        data = np.array(data[2:])
        data.tofile(f"datas/{symbol[1:]}-{t}.bin")
        print(f"[{symbol}-{t}] LOAD OVER {data.shape}")
ed = time.time()
print("total time: ", ed-st)   # 耗时约 4950s

在数据处理过程中，通常需要计算滑动窗口的大小和步长。这两个参数决定了如何从数据中切割出训练样本。滑动窗口的大小定义了每个训练样本的时间窗口长度，而步长定义了滑动窗口之间的间隔, 一旦确定了滑动窗口的大小和步长，接下来会计算每份数据需要从哪些文件中获取数据。这个计算过程通常涉及到迭代数据并根据滑动窗口的设置来确定数据的切割方式。然后，将这些索引信息保存在 index.pkl, 以供后续使用，此阶段耗时约为 4分钟。核心代码如下，详细见 prepare_index.py :

with open("index.pkl", 'wb') as f:
    pickle.dump(index_list, f)
ed = time.time()
print("total time: ", ed-st)    # 约 234s

最后在 Python 代码中，定义一个数据集（DataSet），用于管理和加载训练数据，将 index.pkl 内容读取至内存，使用 mmap 方式打开数据文件，使得能够通过下标访问快速将数据从文件中读取到内存，此阶段耗时约为 4 分钟，核心代码如下，详细见 test_wide_old.py：

def main():
    torch.set_default_device("cuda")
    torch.set_default_tensor_type(torch.DoubleTensor)

    model = SimpleNet()
    model = model.to("cuda")

    loss_fn = nn.MSELoss()
    loss_fn = loss_fn.to("cuda")
    optimizer = torch.optim.Adam(model.parameters(), lr=0.05)
    dataset = MyDataset(4802)
    dataloader = DataLoader(
        dataset, batch_size=64, shuffle=False, num_workers=3,
        pin_memory=True, pin_memory_device="cuda",
        prefetch_factor=5,
    )

    epoch = 10
    for _ in range(epoch):
        for x, y in tqdm(dataloader, mininterval=1):
            x = x.to("cuda")
            y = y.to("cuda")
            y_pred = model(x)
            loss = loss_fn(y_pred, y)
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            
if __name__ == "__main__":
    main()

至此，基于 PytorchDataLoader 深度学习训练数据流程全部结束，第一阶段生成二进制文件大约为 83 分钟，第二阶段生成索引数据信息为 4分钟，第三阶段迭代训练 2 万次耗时为 25 分钟，总耗时为 112 分钟。

在 DDBDataLoader 中，通常会采取以下步骤来处理训练数据：

• 从 DolphinDB 分布式表中加载数据
• 将加载数据处理成训练所需格式

本次测试中，使用了 DDBDataLoader 来获取训练数据。与传统方法不同，无需将数据保存为文件并在客户端进行处理。相反，通过 Session 将 SQL 查询结果直接转换为 torch.Tensor，这可以减少数据传输和存储成本，在测试代码中，使用 Python 中的第三方库 line_profiler 统计各个部分的执行时间，例如数据加载、模型训练等。测试步骤如下：

1. 定义 DDBDataLoader

   在 Python 客户端执行以下代码，使用已建立的数据库表执行 SQL 查询后的结果作为数据集。该数据集指定了目标列为 ["f000001"]，并排除了 Symbol 列和 DateTime 列的数据。此外，还配置了以下参数：

   • batchSize=64 表示 一批数据大小为 64。
   • windowSize=[200, 1], windowStride=[1, 1], offset=200 分别表示输入数据和目标数据的滑动窗口大小分别为 200 和 1，滑动窗口步长分别为1和1，offset为200。
   • shuffle=True 表示数据打乱设置为 True，使用随机种子 seed=0。
   • 使用每支股票的时序数据进行训练，指定 groupCol="Symbol" 和 groupScheme=symbols，其中 symbols 是包含所有股票名称的字符串列表。
   • 为了降低数据分块粒度，指定 repartitionCol="date(DateTime)" 和 repartitionScheme=times，其中 times 是包含 2020.01.01到2020.01.31 所有日期的列表。
   • 训练将在 GPU上进行，指定 device="cuda"，将 torch.Tensor 创建到 GPU 上。
   • prefetchBatch=5, prepartitionNum=3 表示预准备 5 批数据，配置每组查询预载3个子查询的结果。

   这些配置将有助于提高训练效果并充分利用 GPU 和后台线程资源。

   import dolphindb as ddb
   from dolphindb_tools.dataloader import DDBDataLoader
   
   sess = ddb.Session()
   sess.connect('localhost', 8848, "admin", "123456")
   
   dbPath = "dfs://test_ai_dataloader"
   tbName = "wide_factor_table"
   
   symbols = ["`" + f"{i}.SH".zfill(9) for i in range(1, 251)]
   times = ["2020.01." + f"{i+1}".zfill(2) for i in range(31)]
   
   sql = f"""select * from loadTable("{dbPath}", "{tbName}")"""
   
   dataloader = DDBDataLoader(
       s, sql, targetCol=["f000001"], batchSize=64, shuffle=True,
       windowSize=[200, 1], windowStride=[1, 1],
       repartitionCol="date(DateTime)", repartitionScheme=times,
       groupCol="Symbol", groupScheme=symbols,
       seed=0,
       offset=200, excludeCol=["DateTime", "Symbol"], device="cuda",
       prefetchBatch=5, prepartitionNum=3
   )

2. 定义网络并训练

   下述代码在 Python 客户端执行，它定义了一个简单的 CNN 神经网络结构，并定义了损失函数和优化器。最后像使用 torch 中 DataLoader 一样，迭代 DDBDataLoader 获取数据，输入到网络中进行训练，核心代码如下，详细见 test_wide_new.py：

   
   model = SimpleNet()
   model = model.to("cuda")
   loss_fn = nn.MSELoss()
   loss_fn = loss_fn.to("cuda")
   optimizer = torch.optim.Adam(model.parameters(), lr=0.05)
   num_epochs = 10
   
   model.train()
   for epoch in range(num_epochs):
       for X, y in dataloader:
           X = X.to("cuda")
           y = y.to("cuda")
           y_pred = model(X)
           loss = loss_fn(y_pred, y)
           optimizer.zero_grad()
           loss.backward()
           optimizer.step()

   通过将数据直接转换为 torch.Tensor 并使用 DDBDataLoader 管理数据，可以更高效地获取和使用训练数据，从而提高深度学习模型的训练效率。这种方法减少了数据传输和存储的开销，并使训练过程更加灵活和高效。此种方式总耗时为 25 分钟。

从比对结果可以看到，本次测试中，对比了传统方式（PyTorch DataLoader）和 DDBDataLoader，DDBDataLoader 一体化集成 PyTorch 耗时约为 25分钟，内存占用约为 0.8 GB，代码行数约为 70 行， PyTorch DataLoader 总耗时 112 分钟，内存占用约为 4GB，代码行数约为 200 多行。考虑两种方式的特点，原因大概如下：

• 性能提升：在数据准备以及迭代取数耗时方面，DDBDataLoader 耗时明显低于传统方式。这主要是因为 DDBDataLoader 可以直接从 DolphinDB 中直接获得数据，将每个分区（数据源）的数据整体打乱后提供给客户端进行处理。传统方式通常将数据集和 DataLoader 处理逻辑分开，需要先将数据导出成磁盘上的文件，然后再提供给客户端使用，这会对性能产生较大影响。所以，DDBDataLoader 最终相比传统方法有显著性能提升。
• 灵活性增加：在灵活性上，DDBDataLoader 使用 SQL 的方式来初始化，这提供了很高的灵活性。例如，用户可以直接用 SQL 实现新的因子，新实现的因子可以直接应用到 PyTorch 的训练中，而不需要再像传统的方式那样需要将数据再导出成磁盘上的文件。
• 内存占用减少：在内存方面，DolphinDB 内部并行线程以及多消息队列机制，迭代数据集，使用完内存，及时回收，返还给操作系统，减少内存在进程中常驻时间，而传统的数据集和 DataLoader 方式，为直接加载全量数据到内存中，导致内存长时间占用，当涉及数据集过大时，容易产生 OOM 现象。这样 DDBDataLoader 内存使用减少为原来的 1/5。
• 代码行数减少：在代码简洁性方面，DolphinDB 封装了一个 DataLoader 接口，用户使用无感知，只需调用接口，将数据传输到 PyTorch 中，仅仅只需代码 70 行，而 传统的数据集和 DataLoader 需要重新构造一个接口用于数据集与 PyTorch 的对接，代码需 200 多行行。极大的减低了开发运维成本。

综上，DDBDataLoader 可以提升性能以及大幅降低 DolphinDB 内数据用于 PyTorch 训练的开发运维成本。