qlib ai trading

Here’s a hands-on example of implementing a trading strategy with Qlib using Python. This example walks through:

  1. Setting up Qlib
  2. Preparing Data
  3. Defining and Training a Machine Learning Model
  4. Backtesting the Strategy
  5. Evaluating Performance

Step 1: Install and Set Up Qlib

First, install Qlib and required dependencies:

pip install pyqlib

Now, initialize Qlib and download the dataset (we will use the China stock market dataset for demonstration, but you can configure it for other markets):

import qlib
qlib.init(provider_uri="~/.qlib/qlib_data/cn_data")  # Load China market data

Step 2: Load & Prepare Data

Now, let’s prepare our dataset for training:

from qlib.data import D  # Data API
import pandas as pd

# Load market data for a specific stock
df = D.features(["SH600000"], ["$close", "$volume"], start_time="2020-01-01", end_time="2024-12-31")

# Display first few rows
print(df.head())

This will load historical closing prices and trading volume for stock SH600000 (Shanghai Pudong Development Bank).


Step 3: Define & Train an AI Model

Now, we define a machine learning model to predict future stock returns based on historical data:

from qlib.contrib.model.gbdt import LGBModel  # LightGBM Model
from qlib.contrib.strategy.signal_strategy import TopkDropoutStrategy
from qlib.workflow import R  # For Experiment Tracking
from qlib.workflow.record_temp import SignalRecord, PortAnaRecord
from qlib.data.dataset.handler import DataHandlerLP

# Define dataset handler
handler = DataHandlerLP(instruments="csi300", start_time="2010-01-01", end_time="2024-12-31",
                        fit_start_time="2010-01-01", fit_end_time="2018-12-31",
                        infer_start_time="2019-01-01", infer_end_time="2024-12-31")

# Define a LightGBM Model for stock prediction
model = LGBModel()

# Start an experiment
R.start(experiment_name="AI_Trading")

# Train the model
model.fit(handler)

# Save the model
R.save_objects(model=model)
R.end()

Here, we:

  • Use LightGBM (LGB), a powerful ML model for financial data.
  • Train it on historical stock data from the CSI300 index.
  • Save the trained model for backtesting.

Step 4: Backtesting the AI Model

Once trained, we apply the model to make trading decisions and backtest the strategy:

# Load trained model
model = R.load_object("model")

# Generate trading signals
sr = SignalRecord(model=model, dataset=handler, recorder=R.get_recorder())
sr.generate()

# Define a simple trading strategy
strategy = TopkDropoutStrategy(signal="score", topk=50, n_drop=5)

# Backtest the strategy
port_ana = PortAnaRecord(recorder=R.get_recorder(), strategy=strategy)
port_ana.generate()

This performs backtesting using:

  • Signal-based strategy: Picks the top 50 stocks with the highest AI-predicted scores.
  • Dropout mechanism: Avoids excessive turnover by dropping only the bottom 5 stocks from holdings.

Step 5: Evaluate Trading Performance

Finally, let’s visualize the results:

import matplotlib.pyplot as plt

# Load performance results
recorder = R.get_recorder()
df_backtest = recorder.load_object("portfolio_analysis")

# Plot cumulative returns
df_backtest["cumulative_return"].plot(figsize=(10, 5))
plt.title("AI Trading Strategy Performance")
plt.ylabel("Cumulative Return")
plt.xlabel("Date")
plt.show()

This will plot cumulative returns over time, showing how our AI-based strategy performs.


In this example, we:

  • Installed Qlib and prepared stock market data.
  • Trained an AI model (LightGBM) to predict stock returns.
  • Backtested the AI-driven strategy.
  • Visualized performance results.

Deploying AI Trading with Deep Learning (LSTM) in Qlib

In this section, we’ll enhance our AI trading strategy using LSTM (Long Short-Term Memory), a deep learning model well-suited for time-series forecasting. We’ll also discuss how to deploy this strategy in real-time trading.


Step 1: Install Required Libraries

Ensure you have Qlib and PyTorch installed:

pip install pyqlib torch torchvision torchaudio

Step 2: Initialize Qlib and Load Data

We start by initializing Qlib and loading stock market data:

import qlib
from qlib.data import D

qlib.init(provider_uri="~/.qlib/qlib_data/cn_data")  # Use China market data

# Load historical data for a stock
df = D.features(["SH600000"], ["$close", "$volume"], start_time="2023-01-01", end_time="2024-12-31")
print(df.head())  # View the first few rows

Step 3: Define & Train LSTM Model

Now, let’s define an LSTM model using PyTorch:

import torch
import torch.nn as nn
import torch.optim as optim
import numpy as np

# LSTM Model for Stock Prediction
class LSTMModel(nn.Module):
    def __init__(self, input_size, hidden_size, num_layers, output_size):
        super(LSTMModel, self).__init__()
        self.hidden_size = hidden_size
        self.num_layers = num_layers
        self.lstm = nn.LSTM(input_size, hidden_size, num_layers, batch_first=True)
        self.fc = nn.Linear(hidden_size, output_size)

    def forward(self, x):
        h0 = torch.zeros(self.num_layers, x.size(0), self.hidden_size).to(x.device)
        c0 = torch.zeros(self.num_layers, x.size(0), self.hidden_size).to(x.device)
        out, _ = self.lstm(x, (h0, c0))
        out = self.fc(out[:, -1, :])  # Get last time step output
        return out

# Model Parameters
input_size = 2  # Close Price, Volume
hidden_size = 64
num_layers = 2
output_size = 1

# Initialize Model
model = LSTMModel(input_size, hidden_size, num_layers, output_size)
criterion = nn.MSELoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

Step 4: Prepare Data for Training

Convert the stock price data into sequences for LSTM training:

import pandas as pd
from sklearn.preprocessing import MinMaxScaler

# Prepare Data
def create_sequences(data, seq_length):
    xs, ys = [], []
    for i in range(len(data) - seq_length):
        x = data[i:i + seq_length]
        y = data[i + seq_length][0]  # Predict close price
        xs.append(x)
        ys.append(y)
    return np.array(xs), np.array(ys)

# Load Data
df = D.features(["SH600000"], ["$close", "$volume"], start_time="2010-01-01", end_time="2024-12-31")
df.dropna(inplace=True)

# Normalize Data
scaler = MinMaxScaler()
df_scaled = scaler.fit_transform(df)

# Create Sequences
seq_length = 30  # 30 days lookback
X, y = create_sequences(df_scaled, seq_length)

# Convert to PyTorch tensors
X_train = torch.tensor(X[:-500], dtype=torch.float32)  # Train on all but last 500 days
y_train = torch.tensor(y[:-500], dtype=torch.float32).view(-1, 1)
X_test = torch.tensor(X[-500:], dtype=torch.float32)  # Test on last 500 days
y_test = torch.tensor(y[-500:], dtype=torch.float32).view(-1, 1)

Step 5: Train the LSTM Model

Now, train the LSTM model on historical stock data:

num_epochs = 50
for epoch in range(num_epochs):
    model.train()
    optimizer.zero_grad()
    outputs = model(X_train)
    loss = criterion(outputs, y_train)
    loss.backward()
    optimizer.step()
    
    if epoch % 10 == 0:
        print(f"Epoch {epoch}, Loss: {loss.item():.6f}")

print("Training complete!")

Step 6: Backtest AI Strategy Using Qlib

Now, let’s use Qlib to backtest our trained LSTM model: