AI DataLoader

Starting from version 1.30.22.2, DolphinDB Python API provides the deep learning utility class DDBDataLoader , which offers an easy-to-use interface for batch splitting and reshuffling of datasets corresponding to DolphinDB SQL, directly integrating data from DolphinDB into PyTorch.

pip install dolphindb-tools

Expected output:

Note:

• If the result table from the SQL query contains unsupported types, even if the column names are specified in targetCol (which indicates the column names corresponding to y during iteration), those data columns will not appear in the input data or target data.
• The ArrayVector type is supported for the above types. If using ArrayVector columns, you must ensure that all input or target data are of the ArrayVector type.
• torch.bool does not support boolean data with null values. Ensure the BOOL type column contains no nulls before retrieving the data.

The DDBDataLoader class is provided to load and access data, with the interface as follows:

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
)

Server

• Hardware Environment

  Hardware	Configuration
  Host	HostName
  Public IP	xxx.xxx.xxx.218
  Operating System	Linux (Kernel version above 3.10)
  Memory	500 GB
  CPU	x86_64 (64 cores)
  GPU	A100
  Network	10-Gigabit Ethernet

• Software Environment

  Software	Version
  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

• Performance Test Tool

  The third-party Python library line_profiler (version 4.0.3) is used for performance testing. The code to be tested is encapsulated as a function and annotated with the @profile decorator. Then, performance testing is executed via the terminal command kernprof -l -v test.py.

• Test Data

  3-second snapshot frequency factor data is used, generating a total of approximately 277GB of data. The data generation script is executed in the DolphinDB client, specifying Datetime and Symbol as partition and sort columns. A partitioned table wide_factor_table is created in the database dfs://test_ai_dataloader. The table includes a Datetime column, a Symbol column for stock names, and 1000 factor columns (named from f000001 to f001000). The types are DATETIME and SYMBOL, with all factor columns of type DOUBLE. The detailed code is found in the project file ddb_scripts.dos. Core code is as follows:

  dbName = "dfs://test_ai_dataloader"
  tbName = "wide_factor_table"
  
  if (existsDatabase(dbName)) {
      dropDatabase(dbName)
  }
  
  // Number of stocks
  numSymbols = 250
  // Number of factors
  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)

  The following script can be used to check whether the data was successfully written:

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

  Example data is as follows:

  

In traditional deep learning, the following steps are usually used to process training data:

• Generate binary data files.
• Generate index information as a .pkl file.
• Load data into the model using PyTorch's DataLoader.

First, binary data files are generated using the numpy library. This stage takes approximately 83 minutes. Core code is as follows:

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)   # Approx. 4950s

During data processing, it is usually necessary to calculate the sliding window size and stride. These two parameters determine how training samples are sliced from the data. The window size defines the temporal length of each training sample, while the stride defines the gap between successive windows. Once the window size and stride are determined, the next step is to calculate which files each sample needs to load data from. This process typically involves iterating through the data and determining slicing based on the window settings. The index information is then stored in index.pkl for later use. This stage takes approximately 4 minutes. Core code is as follows:

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

Finally, in the Python code, a dataset is defined to manage and load training data. The contents of index.pkl are loaded into memory, and data files are opened using mmap to enable fast access via index. This stage takes approximately 25 minutes. Core code is as follows:

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()

At this point, the PyTorch DataLoader-based deep learning training data pipeline is complete. The first stage (generating binary files) takes approximately 83 minutes. The second stage (generating index information) takes 4 minutes. The third stage (20,000 training iterations) takes 25 minutes. The total time is 112 minutes.

With DDBDataLoader , the following steps are typically used to process training data:

• Load data from the DolphinDB distributed table.
• Process the loaded data into a format required for training.

In this test, DDBDataLoader is used to fetch training data. Unlike traditional methods, it does not require saving data to files for client-side processing. Instead, SQL query results are directly converted into torch.Tensor via a Session, reducing data transmission and storage overhead. In the test code, Python's third-party library line_profiler is used to measure execution time for various parts, such as data loading and model training. The test steps are as follows:

Define DDBDataLoader

The following code is executed in the Python client, using an existing database table. SQL query results serve as the dataset. The dataset specifies the target column as ["f000001"], excluding Symbol and DateTime columns.

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
)

Key parameters explained:

• batchSize=64 indicates the batch size is 64.
• windowSize=[200, 1], windowStride=[1, 1], offset=200 indicates the input and target window sizes are 200 and 1 respectively, with strides of 1, and an offset of 200.
• shuffle=True means data shuffling is enabled, using random seed seed=0
• To train using each stock's time series data, groupCol="Symbol" and groupScheme=symbols are specified, where symbols is a list of all stock names.
• To reduce data partitioning granularity, repartitionCol="date(DateTime)" and repartitionScheme=times are used, where times includes all dates from 2020.01.01 to 2020.01.31.
• Training is performed on GPU by specifying device="cuda", so torch.Tensor is created on the GPU.
• prefetchBatch=5, prepartitionNum=3 means 5 batches are prepared in advance, and 3 subquery results are preloaded per partition.

These configurations help improve training efficiency and fully utilize GPU and background threads.

Define Network and Train

The following code, executed in the Python client, defines a simple CNN network structure, along with the loss function and optimizer. Finally, similar to torch’s DataLoader, it iterates through the DDBDataLoader to retrieve and input data into the model for training. Core code is as follows:

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()

By directly converting data to torch.Tensor and using DDBDataLoader for data management, training data can be accessed and utilized more efficiently, improving deep learning model training efficiency. This approach reduces data transmission and storage overhead, and makes the training process more flexible and efficient. Total time cost for this method is 25 minutes.

Performance Comparison

This test compares the traditional method (PyTorch DataLoader) and DDBDataLoader. DDBDataLoader, with integrated PyTorch support, takes about 25 minutes, uses 0.8 GB memory, and 70 lines of code. In contrast, PyTorch DataLoader takes 112 minutes, 4 GB memory, and over 200 lines of code. Key reasons include:

• Performance Improvement: DDBDataLoader significantly reduces data preparation and iteration time. It retrieves and shuffles data partitions directly from DolphinDB, avoiding the traditional export-to-disk and re-import workflow that hinders performance.
• Increased Flexibility: DDBDataLoader initializes using SQL, allowing high flexibility. Users can create new factors via SQL and immediately apply them in PyTorch training—unlike traditional methods, which require re-exporting data.
• Reduced Memory Usage: DolphinDB employs internal parallel threads and message queues, recycling memory after use. Traditional methods often load full datasets into memory, leading to long memory residency and risk of OOM. DDBDataLoader reduces memory usage to 1/5 of the traditional method.
• Reduced Code Complexity: DolphinDB wraps a DataLoader interface with minimal user intervention—only 70 lines of code. The traditional approach needs custom interfaces for dataset and PyTorch integration, totaling over 200 lines. This greatly reduces development and maintenance costs.

In summary, DDBDataLoader can improve performance and significantly reduce development and operational costs when training PyTorch models with data stored in DolphinDB.