Home

Awesome

Stock Trend Prediction with Multi-granularity Data: A Contrastive Learning Approach with Adaptive Fusion

  1. Dependencies

    Install packages from requirements.txt.

  2. Load Data using qlib

    $ cd ./load_data
    

    Download daily data:

    $ python load_dataset.py
    
    • Change parameter market to get data from different dataset: csi300, csi800, NASDAQ etc.

      Data Sample - SH600000 in CSI300

      features dimensions = 6 * 20 + 1 = 121

    Download high-frequency data:

    $ python high_freq_resample.py
    
    • Change parameter N to get data from different frequencies: 15min, 30min, 120min etc.

      Data Sample - SH600000 in CSI300

      features dimensions = 16 * 6 * 20 + 1 = 1921

  3. Framework

  1. Run

$ cd ./framework

Train Pre-train model:

$ python main_contrast.py with config/contrast_all_2_encoder.json model_name=contrastive_all_2_encoder

Train Adaptive Multi-granularity Feature Fusion model:

$ python main_contrast_2_stage.py with config/contrast_all_2_stage.json model_name=contrastive_all_2_stage

Run Market Trading Simulation:

$ cd ./framework
$ python trade_sim.py
  1. Records

    Records for each experiment are saved in ./framework/my_runs/.
    Each record file includes:

    config.json

    • contains the parameter settings and data path.

    cout.txt

    • contains the name of dataset, detailed model output, and experiment results.

    pred_{model_name}_{seed}.pkl

>  * contains the  `score` (model prediction) and `label`
> run.json

* contains the hash ids of every script used in the experiment. And the source code can be found in `./framework/my_runs/source/`.