From 412c9eee2ef333fb596220e17576384b83655aba Mon Sep 17 00:00:00 2001 From: lwwang1995 Date: Sun, 6 Dec 2020 22:24:03 +0800 Subject: [PATCH] Update models. --- .../ALSTM/workflow_config_alstm_Alpha158.yaml | 93 ++++ .../workflow_config_catboost_Alpha360.yaml | 72 +++ .../GATs/workflow_config_gats_Alpha158.yaml | 92 ++++ .../GRU/workflow_config_gru_Alpha158.yaml | 93 ++++ .../LSTM/workflow_config_lstm_Alpha158.yaml | 93 ++++ .../workflow_config_lightgbm_Alpha360.yaml | 73 +++ .../MLP/workflow_config_mlp_Alpha360.yaml | 82 ++++ .../SFM/workflow_config_alstm_Alpha158.yaml | 93 ++++ .../workflow_config_xgboost_Alpha360.yaml | 71 +++ qlib/contrib/model/pytorch_alstm_ts.py | 331 +++++++++++++ qlib/contrib/model/pytorch_gats_ts.py | 379 +++++++++++++++ qlib/contrib/model/pytorch_gru_ts.py | 302 ++++++++++++ qlib/contrib/model/pytorch_lstm_ts.py | 303 ++++++++++++ qlib/contrib/model/pytorch_sfm_ts.py | 434 ++++++++++++++++++ 14 files changed, 2511 insertions(+) create mode 100755 examples/benchmarks/ALSTM/workflow_config_alstm_Alpha158.yaml create mode 100644 examples/benchmarks/CatBoost/workflow_config_catboost_Alpha360.yaml create mode 100644 examples/benchmarks/GATs/workflow_config_gats_Alpha158.yaml create mode 100755 examples/benchmarks/GRU/workflow_config_gru_Alpha158.yaml create mode 100755 examples/benchmarks/LSTM/workflow_config_lstm_Alpha158.yaml create mode 100644 examples/benchmarks/LightGBM/workflow_config_lightgbm_Alpha360.yaml create mode 100644 examples/benchmarks/MLP/workflow_config_mlp_Alpha360.yaml create mode 100755 examples/benchmarks/SFM/workflow_config_alstm_Alpha158.yaml create mode 100644 examples/benchmarks/XGBoost/workflow_config_xgboost_Alpha360.yaml create mode 100644 qlib/contrib/model/pytorch_alstm_ts.py create mode 100644 qlib/contrib/model/pytorch_gats_ts.py create mode 100755 qlib/contrib/model/pytorch_gru_ts.py create mode 100755 qlib/contrib/model/pytorch_lstm_ts.py create mode 100644 qlib/contrib/model/pytorch_sfm_ts.py diff --git a/examples/benchmarks/ALSTM/workflow_config_alstm_Alpha158.yaml b/examples/benchmarks/ALSTM/workflow_config_alstm_Alpha158.yaml new file mode 100755 index 000000000..878e3c065 --- /dev/null +++ b/examples/benchmarks/ALSTM/workflow_config_alstm_Alpha158.yaml @@ -0,0 +1,93 @@ +qlib_init: + provider_uri: "~/.qlib/qlib_data/cn_data" + region: cn +market: &market csi300 +benchmark: &benchmark SH000300 +data_handler_config: &data_handler_config + start_time: 2008-01-01 + end_time: 2020-08-01 + fit_start_time: 2008-01-01 + fit_end_time: 2014-12-31 + instruments: *market + infer_processors: + - class: FilterCol + kwargs: + fields_group: feature + col_list: ["RESI5", "WVMA5", "RSQR5", "KLEN", "RSQR10", "CORR5", "CORD5", "CORR10", + "ROC60", "RESI10", "VSTD5", "RSQR60", "CORR60", "WVMA60", "STD5", + "RSQR20", "CORD60", "CORD10", "CORR20", "KLOW" + ] + - class: RobustZScoreNorm + kwargs: + fields_group: feature + clip_outlier: true + - class: Fillna + kwargs: + fields_group: feature + learn_processors: + - class: DropnaLabel + - class: CSRankNorm + kwargs: + fields_group: label + label: ["Ref($close, -2) / Ref($close, -1) - 1"] + +port_analysis_config: &port_analysis_config + strategy: + class: TopkDropoutStrategy + module_path: qlib.contrib.strategy.strategy + kwargs: + topk: 50 + n_drop: 5 + backtest: + verbose: False + limit_threshold: 0.095 + account: 100000000 + benchmark: *benchmark + deal_price: close + open_cost: 0.0005 + close_cost: 0.0015 + min_cost: 5 +task: + model: + class: ALSTM + module_path: qlib.contrib.model.pytorch_alstm_ts + kwargs: + d_feat: 20 + hidden_size: 64 + num_layers: 2 + dropout: 0.0 + n_epochs: 200 + lr: 1e-3 + early_stop: 10 + batch_size: 800 + metric: loss + loss: mse + n_jobs: 20 + GPU: 0 + rnn_type: GRU + dataset: + class: TSDatasetH + module_path: qlib.data.dataset + kwargs: + handler: + class: Alpha158 + module_path: qlib.contrib.data.handler + kwargs: *data_handler_config + segments: + train: [2008-01-01, 2014-12-31] + valid: [2015-01-01, 2016-12-31] + test: [2017-01-01, 2020-08-01] + step_len: 20 + record: + - class: SignalRecord + module_path: qlib.workflow.record_temp + kwargs: {} + - class: SigAnaRecord + module_path: qlib.workflow.record_temp + kwargs: + ana_long_short: False + ann_scaler: 252 + - class: PortAnaRecord + module_path: qlib.workflow.record_temp + kwargs: + config: *port_analysis_config diff --git a/examples/benchmarks/CatBoost/workflow_config_catboost_Alpha360.yaml b/examples/benchmarks/CatBoost/workflow_config_catboost_Alpha360.yaml new file mode 100644 index 000000000..f7dc26f5d --- /dev/null +++ b/examples/benchmarks/CatBoost/workflow_config_catboost_Alpha360.yaml @@ -0,0 +1,72 @@ +qlib_init: + provider_uri: "~/.qlib/qlib_data/cn_data" + region: cn +market: &market csi300 +benchmark: &benchmark SH000300 +data_handler_config: &data_handler_config + start_time: 2008-01-01 + end_time: 2020-08-01 + fit_start_time: 2008-01-01 + fit_end_time: 2014-12-31 + instruments: *market + infer_processors: [] + learn_processors: + - class: DropnaLabel + - class: CSRankNorm + kwargs: + fields_group: label + label: ["Ref($close, -2) / Ref($close, -1) - 1"] +port_analysis_config: &port_analysis_config + strategy: + class: TopkDropoutStrategy + module_path: qlib.contrib.strategy.strategy + kwargs: + topk: 50 + n_drop: 5 + backtest: + verbose: False + limit_threshold: 0.095 + account: 100000000 + benchmark: *benchmark + deal_price: close + open_cost: 0.0005 + close_cost: 0.0015 + min_cost: 5 +task: + model: + class: CatBoostModel + module_path: qlib.contrib.model.catboost_model + kwargs: + loss: RMSE + learning_rate: 0.0421 + subsample: 0.8789 + max_depth: 6 + num_leaves: 100 + thread_count: 20 + grow_policy: Lossguide + bootstrap_type: Poisson + dataset: + class: DatasetH + module_path: qlib.data.dataset + kwargs: + handler: + class: Alpha360 + module_path: qlib.contrib.data.handler + kwargs: *data_handler_config + segments: + train: [2008-01-01, 2014-12-31] + valid: [2015-01-01, 2016-12-31] + test: [2017-01-01, 2020-08-01] + record: + - class: SignalRecord + module_path: qlib.workflow.record_temp + kwargs: {} + - class: SigAnaRecord + module_path: qlib.workflow.record_temp + kwargs: + ana_long_short: False + ann_scaler: 252 + - class: PortAnaRecord + module_path: qlib.workflow.record_temp + kwargs: + config: *port_analysis_config diff --git a/examples/benchmarks/GATs/workflow_config_gats_Alpha158.yaml b/examples/benchmarks/GATs/workflow_config_gats_Alpha158.yaml new file mode 100644 index 000000000..c0cb8338a --- /dev/null +++ b/examples/benchmarks/GATs/workflow_config_gats_Alpha158.yaml @@ -0,0 +1,92 @@ +qlib_init: + provider_uri: "~/.qlib/qlib_data/cn_data" + region: cn +market: &market csi300 +benchmark: &benchmark SH000300 +data_handler_config: &data_handler_config + start_time: 2008-01-01 + end_time: 2020-08-01 + fit_start_time: 2008-01-01 + fit_end_time: 2014-12-31 + instruments: *market + infer_processors: + - class: FilterCol + kwargs: + fields_group: feature + col_list: ["RESI5", "WVMA5", "RSQR5", "KLEN", "RSQR10", "CORR5", "CORD5", "CORR10", + "ROC60", "RESI10", "VSTD5", "RSQR60", "CORR60", "WVMA60", "STD5", + "RSQR20", "CORD60", "CORD10", "CORR20", "KLOW" + ] + - class: RobustZScoreNorm + kwargs: + fields_group: feature + clip_outlier: true + - class: Fillna + kwargs: + fields_group: feature + learn_processors: + - class: DropnaLabel + - class: CSRankNorm + kwargs: + fields_group: label + label: ["Ref($close, -2) / Ref($close, -1) - 1"] +port_analysis_config: &port_analysis_config + strategy: + class: TopkDropoutStrategy + module_path: qlib.contrib.strategy.strategy + kwargs: + topk: 50 + n_drop: 5 + backtest: + verbose: False + limit_threshold: 0.095 + account: 100000000 + benchmark: *benchmark + deal_price: close + open_cost: 0.0005 + close_cost: 0.0015 + min_cost: 5 +task: + model: + class: GATs + module_path: qlib.contrib.model.pytorch_gats_ts + kwargs: + d_feat: 6 + hidden_size: 64 + num_layers: 2 + dropout: 0.7 + n_epochs: 200 + lr: 1e-4 + early_stop: 20 + metric: loss + loss: mse + base_model: LSTM + with_pretrain: True + model_path: "benchmarks/LSTM/model_lstm_ts.pkl" + GPU: 0 + dataset: + class: TSDatasetH + module_path: qlib.data.dataset + kwargs: + handler: + class: Alpha158 + module_path: qlib.contrib.data.handler + kwargs: *data_handler_config + segments: + train: [2008-01-01, 2014-12-31] + valid: [2015-01-01, 2016-12-31] + test: [2017-01-01, 2020-08-01] + step_len: 20 + record: + - class: SignalRecord + module_path: qlib.workflow.record_temp + kwargs: {} + - class: SigAnaRecord + module_path: qlib.workflow.record_temp + kwargs: + ana_long_short: False + ann_scaler: 252 + - class: PortAnaRecord + module_path: qlib.workflow.record_temp + kwargs: + config: *port_analysis_config \ No newline at end of file diff --git a/examples/benchmarks/GRU/workflow_config_gru_Alpha158.yaml b/examples/benchmarks/GRU/workflow_config_gru_Alpha158.yaml new file mode 100755 index 000000000..61567e5fc --- /dev/null +++ b/examples/benchmarks/GRU/workflow_config_gru_Alpha158.yaml @@ -0,0 +1,93 @@ +qlib_init: + provider_uri: "~/.qlib/qlib_data/cn_data" + region: cn +market: &market csi300 +benchmark: &benchmark SH000300 +data_handler_config: &data_handler_config + start_time: 2008-01-01 + end_time: 2020-08-01 + fit_start_time: 2008-01-01 + fit_end_time: 2014-12-31 + instruments: *market + infer_processors: + - class: FilterCol + kwargs: + fields_group: feature + col_list: ["RESI5", "WVMA5", "RSQR5", "KLEN", "RSQR10", "CORR5", "CORD5", "CORR10", + "ROC60", "RESI10", "VSTD5", "RSQR60", "CORR60", "WVMA60", "STD5", + "RSQR20", "CORD60", "CORD10", "CORR20", "KLOW" + ] + - class: RobustZScoreNorm + kwargs: + fields_group: feature + clip_outlier: true + - class: Fillna + kwargs: + fields_group: feature + learn_processors: + - class: DropnaLabel + - class: CSRankNorm + kwargs: + fields_group: label + label: ["Ref($close, -2) / Ref($close, -1) - 1"] + +port_analysis_config: &port_analysis_config + strategy: + class: TopkDropoutStrategy + module_path: qlib.contrib.strategy.strategy + kwargs: + topk: 50 + n_drop: 5 + backtest: + verbose: False + limit_threshold: 0.095 + account: 100000000 + benchmark: *benchmark + deal_price: close + open_cost: 0.0005 + close_cost: 0.0015 + min_cost: 5 +task: + model: + class: GRU + module_path: qlib.contrib.model.pytorch_gru_ts + kwargs: + d_feat: 20 + hidden_size: 64 + num_layers: 2 + dropout: 0.0 + n_epochs: 2 + lr: 1e-3 + early_stop: 10 + batch_size: 800 + metric: loss + loss: mse + n_jobs: 20 + GPU: 0 + rnn_type: GRU + dataset: + class: TSDatasetH + module_path: qlib.data.dataset + kwargs: + handler: + class: Alpha158 + module_path: qlib.contrib.data.handler + kwargs: *data_handler_config + segments: + train: [2008-01-01, 2014-12-31] + valid: [2015-01-01, 2016-12-31] + test: [2017-01-01, 2020-08-01] + step_len: 20 + record: + - class: SignalRecord + module_path: qlib.workflow.record_temp + kwargs: {} + - class: SigAnaRecord + module_path: qlib.workflow.record_temp + kwargs: + ana_long_short: False + ann_scaler: 252 + - class: PortAnaRecord + module_path: qlib.workflow.record_temp + kwargs: + config: *port_analysis_config diff --git a/examples/benchmarks/LSTM/workflow_config_lstm_Alpha158.yaml b/examples/benchmarks/LSTM/workflow_config_lstm_Alpha158.yaml new file mode 100755 index 000000000..eedff35a6 --- /dev/null +++ b/examples/benchmarks/LSTM/workflow_config_lstm_Alpha158.yaml @@ -0,0 +1,93 @@ +qlib_init: + provider_uri: "~/.qlib/qlib_data/cn_data" + region: cn +market: &market csi300 +benchmark: &benchmark SH000300 +data_handler_config: &data_handler_config + start_time: 2008-01-01 + end_time: 2020-08-01 + fit_start_time: 2008-01-01 + fit_end_time: 2014-12-31 + instruments: *market + infer_processors: + - class: FilterCol + kwargs: + fields_group: feature + col_list: ["RESI5", "WVMA5", "RSQR5", "KLEN", "RSQR10", "CORR5", "CORD5", "CORR10", + "ROC60", "RESI10", "VSTD5", "RSQR60", "CORR60", "WVMA60", "STD5", + "RSQR20", "CORD60", "CORD10", "CORR20", "KLOW" + ] + - class: RobustZScoreNorm + kwargs: + fields_group: feature + clip_outlier: true + - class: Fillna + kwargs: + fields_group: feature + learn_processors: + - class: DropnaLabel + - class: CSRankNorm + kwargs: + fields_group: label + label: ["Ref($close, -2) / Ref($close, -1) - 1"] + +port_analysis_config: &port_analysis_config + strategy: + class: TopkDropoutStrategy + module_path: qlib.contrib.strategy.strategy + kwargs: + topk: 50 + n_drop: 5 + backtest: + verbose: False + limit_threshold: 0.095 + account: 100000000 + benchmark: *benchmark + deal_price: close + open_cost: 0.0005 + close_cost: 0.0015 + min_cost: 5 +task: + model: + class: LSTM + module_path: qlib.contrib.model.pytorch_lstm_ts + kwargs: + d_feat: 20 + hidden_size: 64 + num_layers: 2 + dropout: 0.0 + n_epochs: 200 + lr: 1e-3 + early_stop: 10 + batch_size: 800 + metric: loss + loss: mse + n_jobs: 20 + GPU: 0 + rnn_type: GRU + dataset: + class: TSDatasetH + module_path: qlib.data.dataset + kwargs: + handler: + class: Alpha158 + module_path: qlib.contrib.data.handler + kwargs: *data_handler_config + segments: + train: [2008-01-01, 2014-12-31] + valid: [2015-01-01, 2016-12-31] + test: [2017-01-01, 2020-08-01] + step_len: 20 + record: + - class: SignalRecord + module_path: qlib.workflow.record_temp + kwargs: {} + - class: SigAnaRecord + module_path: qlib.workflow.record_temp + kwargs: + ana_long_short: False + ann_scaler: 252 + - class: PortAnaRecord + module_path: qlib.workflow.record_temp + kwargs: + config: *port_analysis_config diff --git a/examples/benchmarks/LightGBM/workflow_config_lightgbm_Alpha360.yaml b/examples/benchmarks/LightGBM/workflow_config_lightgbm_Alpha360.yaml new file mode 100644 index 000000000..92df15133 --- /dev/null +++ b/examples/benchmarks/LightGBM/workflow_config_lightgbm_Alpha360.yaml @@ -0,0 +1,73 @@ +qlib_init: + provider_uri: "~/.qlib/qlib_data/cn_data" + region: cn +market: &market csi300 +benchmark: &benchmark SH000300 +data_handler_config: &data_handler_config + start_time: 2008-01-01 + end_time: 2020-08-01 + fit_start_time: 2008-01-01 + fit_end_time: 2014-12-31 + instruments: *market + infer_processors: [] + learn_processors: + - class: DropnaLabel + - class: CSRankNorm + kwargs: + fields_group: label + label: ["Ref($close, -2) / Ref($close, -1) - 1"] +port_analysis_config: &port_analysis_config + strategy: + class: TopkDropoutStrategy + module_path: qlib.contrib.strategy.strategy + kwargs: + topk: 50 + n_drop: 5 + backtest: + verbose: False + limit_threshold: 0.095 + account: 100000000 + benchmark: *benchmark + deal_price: close + open_cost: 0.0005 + close_cost: 0.0015 + min_cost: 5 +task: + model: + class: LGBModel + module_path: qlib.contrib.model.gbdt + kwargs: + loss: mse + colsample_bytree: 0.8879 + learning_rate: 0.0421 + subsample: 0.8789 + lambda_l1: 205.6999 + lambda_l2: 580.9768 + max_depth: 8 + num_leaves: 210 + num_threads: 20 + dataset: + class: DatasetH + module_path: qlib.data.dataset + kwargs: + handler: + class: Alpha360 + module_path: qlib.contrib.data.handler + kwargs: *data_handler_config + segments: + train: [2008-01-01, 2014-12-31] + valid: [2015-01-01, 2016-12-31] + test: [2017-01-01, 2020-08-01] + record: + - class: SignalRecord + module_path: qlib.workflow.record_temp + kwargs: {} + - class: SigAnaRecord + module_path: qlib.workflow.record_temp + kwargs: + ana_long_short: False + ann_scaler: 252 + - class: PortAnaRecord + module_path: qlib.workflow.record_temp + kwargs: + config: *port_analysis_config \ No newline at end of file diff --git a/examples/benchmarks/MLP/workflow_config_mlp_Alpha360.yaml b/examples/benchmarks/MLP/workflow_config_mlp_Alpha360.yaml new file mode 100644 index 000000000..18920399f --- /dev/null +++ b/examples/benchmarks/MLP/workflow_config_mlp_Alpha360.yaml @@ -0,0 +1,82 @@ +qlib_init: + provider_uri: "~/.qlib/qlib_data/cn_data" + region: cn +market: &market csi300 +benchmark: &benchmark SH000300 +data_handler_config: &data_handler_config + start_time: 2008-01-01 + end_time: 2020-08-01 + fit_start_time: 2008-01-01 + fit_end_time: 2014-12-31 + instruments: *market + infer_processors: + - class: RobustZScoreNorm + kwargs: + fields_group: feature + clip_outlier: true + - class: Fillna + kwargs: + fields_group: feature + learn_processors: + - class: DropnaLabel + - class: CSRankNorm + kwargs: + fields_group: label + label: ["Ref($close, -2) / Ref($close, -1) - 1"] + +port_analysis_config: &port_analysis_config + strategy: + class: TopkDropoutStrategy + module_path: qlib.contrib.strategy.strategy + kwargs: + topk: 50 + n_drop: 5 + backtest: + verbose: False + limit_threshold: 0.095 + account: 100000000 + benchmark: *benchmark + deal_price: close + open_cost: 0.0005 + close_cost: 0.0015 + min_cost: 5 +task: + model: + class: DNNModelPytorch + module_path: qlib.contrib.model.pytorch_nn + kwargs: + loss: mse + input_dim: 360 + output_dim: 1 + lr: 0.002 + lr_decay: 0.96 + lr_decay_steps: 100 + optimizer: adam + max_steps: 8000 + batch_size: 4096 + GPU: 0 + dataset: + class: DatasetH + module_path: qlib.data.dataset + kwargs: + handler: + class: Alpha360 + module_path: qlib.contrib.data.handler + kwargs: *data_handler_config + segments: + train: [2008-01-01, 2014-12-31] + valid: [2015-01-01, 2016-12-31] + test: [2017-01-01, 2020-08-01] + record: + - class: SignalRecord + module_path: qlib.workflow.record_temp + kwargs: {} + - class: SigAnaRecord + module_path: qlib.workflow.record_temp + kwargs: + ana_long_short: False + ann_scaler: 252 + - class: PortAnaRecord + module_path: qlib.workflow.record_temp + kwargs: + config: *port_analysis_config \ No newline at end of file diff --git a/examples/benchmarks/SFM/workflow_config_alstm_Alpha158.yaml b/examples/benchmarks/SFM/workflow_config_alstm_Alpha158.yaml new file mode 100755 index 000000000..f196cb086 --- /dev/null +++ b/examples/benchmarks/SFM/workflow_config_alstm_Alpha158.yaml @@ -0,0 +1,93 @@ +qlib_init: + provider_uri: "~/.qlib/qlib_data/cn_data" + region: cn +market: &market csi300 +benchmark: &benchmark SH000300 +data_handler_config: &data_handler_config + start_time: 2008-01-01 + end_time: 2020-08-01 + fit_start_time: 2008-01-01 + fit_end_time: 2014-12-31 + instruments: *market + infer_processors: + - class: FilterCol + kwargs: + fields_group: feature + col_list: ["RESI5", "WVMA5", "RSQR5", "KLEN", "RSQR10", "CORR5", "CORD5", "CORR10", + "ROC60", "RESI10", "VSTD5", "RSQR60", "CORR60", "WVMA60", "STD5", + "RSQR20", "CORD60", "CORD10", "CORR20", "KLOW" + ] + - class: RobustZScoreNorm + kwargs: + fields_group: feature + clip_outlier: true + - class: Fillna + kwargs: + fields_group: feature + learn_processors: + - class: DropnaLabel + - class: CSRankNorm + kwargs: + fields_group: label + label: ["Ref($close, -2) / Ref($close, -1) - 1"] + +port_analysis_config: &port_analysis_config + strategy: + class: TopkDropoutStrategy + module_path: qlib.contrib.strategy.strategy + kwargs: + topk: 50 + n_drop: 5 + backtest: + verbose: False + limit_threshold: 0.095 + account: 100000000 + benchmark: *benchmark + deal_price: close + open_cost: 0.0005 + close_cost: 0.0015 + min_cost: 5 +task: + model: + class: SFM + module_path: qlib.contrib.model.pytorch_sfm_ts + kwargs: + d_feat: 20 + hidden_size: 64 + num_layers: 2 + dropout: 0.0 + n_epochs: 200 + lr: 1e-3 + early_stop: 10 + batch_size: 800 + metric: loss + loss: mse + n_jobs: 20 + GPU: 0 + rnn_type: GRU + dataset: + class: TSDatasetH + module_path: qlib.data.dataset + kwargs: + handler: + class: Alpha158 + module_path: qlib.contrib.data.handler + kwargs: *data_handler_config + segments: + train: [2008-01-01, 2014-12-31] + valid: [2015-01-01, 2016-12-31] + test: [2017-01-01, 2020-08-01] + step_len: 20 + record: + - class: SignalRecord + module_path: qlib.workflow.record_temp + kwargs: {} + - class: SigAnaRecord + module_path: qlib.workflow.record_temp + kwargs: + ana_long_short: False + ann_scaler: 252 + - class: PortAnaRecord + module_path: qlib.workflow.record_temp + kwargs: + config: *port_analysis_config diff --git a/examples/benchmarks/XGBoost/workflow_config_xgboost_Alpha360.yaml b/examples/benchmarks/XGBoost/workflow_config_xgboost_Alpha360.yaml new file mode 100644 index 000000000..7887a25a6 --- /dev/null +++ b/examples/benchmarks/XGBoost/workflow_config_xgboost_Alpha360.yaml @@ -0,0 +1,71 @@ +qlib_init: + provider_uri: "~/.qlib/qlib_data/cn_data" + region: cn +market: &market csi300 +benchmark: &benchmark SH000300 +data_handler_config: &data_handler_config + start_time: 2008-01-01 + end_time: 2020-08-01 + fit_start_time: 2008-01-01 + fit_end_time: 2014-12-31 + instruments: *market + infer_processors: [] + learn_processors: + - class: DropnaLabel + - class: CSRankNorm + kwargs: + fields_group: label + label: ["Ref($close, -2) / Ref($close, -1) - 1"] +port_analysis_config: &port_analysis_config + strategy: + class: TopkDropoutStrategy + module_path: qlib.contrib.strategy.strategy + kwargs: + topk: 50 + n_drop: 5 + backtest: + verbose: False + limit_threshold: 0.095 + account: 100000000 + benchmark: *benchmark + deal_price: close + open_cost: 0.0005 + close_cost: 0.0015 + min_cost: 5 +task: + model: + class: XGBModel + module_path: qlib.contrib.model.xgboost + kwargs: + eval_metric: rmse + colsample_bytree: 0.8879 + eta: 0.0421 + max_depth: 8 + n_estimators: 647 + subsample: 0.8789 + nthread: 20 + dataset: + class: DatasetH + module_path: qlib.data.dataset + kwargs: + handler: + class: Alpha360 + module_path: qlib.contrib.data.handler + kwargs: *data_handler_config + segments: + train: [2008-01-01, 2014-12-31] + valid: [2015-01-01, 2016-12-31] + test: [2017-01-01, 2020-08-01] + record: + - class: SignalRecord + module_path: qlib.workflow.record_temp + kwargs: {} + - class: SigAnaRecord + module_path: qlib.workflow.record_temp + kwargs: + ana_long_short: False + ann_scaler: 252 + - class: PortAnaRecord + module_path: qlib.workflow.record_temp + kwargs: + config: *port_analysis_config diff --git a/qlib/contrib/model/pytorch_alstm_ts.py b/qlib/contrib/model/pytorch_alstm_ts.py new file mode 100644 index 000000000..411a60967 --- /dev/null +++ b/qlib/contrib/model/pytorch_alstm_ts.py @@ -0,0 +1,331 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + + +from __future__ import division +from __future__ import print_function + +import os +import numpy as np +import pandas as pd +import copy +from sklearn.metrics import roc_auc_score, mean_squared_error +import logging +from ...utils import ( + unpack_archive_with_buffer, + save_multiple_parts_file, + create_save_path, + drop_nan_by_y_index, +) +from ...log import get_module_logger, TimeInspector + +import torch +import torch.nn as nn +import torch.optim as optim +from torch.utils.data import DataLoader + +from ...model.base import Model +from ...data.dataset import DatasetH, TSDatasetH +from ...data.dataset.handler import DataHandlerLP + + +class ALSTM(Model): + """ALSTM Model + + Parameters + ---------- + d_feat : int + input dimension for each time step + metric: str + the evaluate metric used in early stop + optimizer : str + optimizer name + GPU : str + the GPU ID(s) used for training + """ + + def __init__( + self, + d_feat=6, + hidden_size=64, + num_layers=2, + dropout=0.0, + n_epochs=200, + lr=0.001, + metric="", + batch_size=2000, + early_stop=20, + loss="mse", + optimizer="adam", + n_jobs=10, + GPU="0", + seed=None, + **kwargs + ): + # Set logger. + self.logger = get_module_logger("ALSTM") + self.logger.info("ALSTM pytorch version...") + + # set hyper-parameters. + self.d_feat = d_feat + self.hidden_size = hidden_size + self.num_layers = num_layers + self.dropout = dropout + self.n_epochs = n_epochs + self.lr = lr + self.metric = metric + self.batch_size = batch_size + self.early_stop = early_stop + self.optimizer = optimizer.lower() + self.loss = loss + self.device = torch.device("cuda:%d" % (GPU) if torch.cuda.is_available() else "cpu") + self.n_jobs = n_jobs + self.use_gpu = torch.cuda.is_available() + self.seed = seed + + self.logger.info( + "ALSTM parameters setting:" + "\nd_feat : {}" + "\nhidden_size : {}" + "\nnum_layers : {}" + "\ndropout : {}" + "\nn_epochs : {}" + "\nlr : {}" + "\nmetric : {}" + "\nbatch_size : {}" + "\nearly_stop : {}" + "\noptimizer : {}" + "\nloss_type : {}" + "\nvisible_GPU : {}" + "\nn_jobs : {}" + "\nuse_GPU : {}" + "\nseed : {}".format( + d_feat, + hidden_size, + num_layers, + dropout, + n_epochs, + lr, + metric, + batch_size, + early_stop, + optimizer.lower(), + loss, + GPU, + n_jobs, + self.use_gpu, + seed, + ) + ) + + if self.seed is not None: + np.random.seed(self.seed) + torch.manual_seed(self.seed) + + self.ALSTM_model = ALSTMModel( + d_feat=self.d_feat, + hidden_size=self.hidden_size, + num_layers=self.num_layers, + dropout=self.dropout, + ).to(self.device) + if optimizer.lower() == "adam": + self.train_optimizer = optim.Adam(self.ALSTM_model.parameters(), lr=self.lr) + elif optimizer.lower() == "gd": + self.train_optimizer = optim.SGD(self.ALSTM_model.parameters(), lr=self.lr) + else: + raise NotImplementedError("optimizer {} is not supported!".format(optimizer)) + + self._fitted = False + self.ALSTM_model.to(self.device) + + def mse(self, pred, label): + loss = (pred - label) ** 2 + return torch.mean(loss) + + def loss_fn(self, pred, label): + mask = ~torch.isnan(label) + + if self.loss == "mse": + return self.mse(pred[mask], label[mask]) + + raise ValueError("unknown loss `%s`" % self.loss) + + def metric_fn(self, pred, label): + + mask = torch.isfinite(label) + + if self.metric == "" or self.metric == "loss": + return -self.loss_fn(pred[mask], label[mask]) + + raise ValueError("unknown metric `%s`" % self.metric) + + def train_epoch(self, data_loader): + + self.ALSTM_model.train() + + for data in data_loader: + feature = data[:,:,0:-1].to(self.device) + label = data[:,-1,-1].to(self.device) + + pred = self.ALSTM_model(feature.float()) + loss = self.loss_fn(pred, label) + + self.train_optimizer.zero_grad() + loss.backward() + torch.nn.utils.clip_grad_value_(self.ALSTM_model.parameters(), 3.0) + self.train_optimizer.step() + + def test_epoch(self, data_loader): + + self.ALSTM_model.eval() + + scores = [] + losses = [] + + for data in data_loader: + + feature = data[:,:,0:-1].to(self.device) + # feature[torch.isnan(feature)] = 0 + label = data[:,-1,-1].to(self.device) + + pred = self.ALSTM_model(feature.float()) + loss = self.loss_fn(pred, label) + losses.append(loss.item()) + + score = self.metric_fn(pred, label) + scores.append(score.item()) + + return np.mean(losses), np.mean(scores) + + def fit( + self, + dataset, + evals_result=dict(), + verbose=True, + save_path=None, + ): + dl_train = dataset.prepare("train", data_key=DataHandlerLP.DK_L) + dl_valid = dataset.prepare("valid", data_key=DataHandlerLP.DK_L) + + dl_train.config(fillna_type="ffill+bfill") # process nan brought by dataloader + dl_valid.config(fillna_type="ffill+bfill") # process nan brought by dataloader + + train_loader = DataLoader(dl_train, batch_size=self.batch_size, shuffle=True, num_workers=self.n_jobs) + valid_loader = DataLoader(dl_valid, batch_size=self.batch_size, shuffle=False, num_workers=self.n_jobs) + + if save_path == None: + save_path = create_save_path(save_path) + + stop_steps = 0 + train_loss = 0 + best_score = -np.inf + best_epoch = 0 + evals_result["train"] = [] + evals_result["valid"] = [] + + # train + self.logger.info("training...") + self._fitted = True + + for step in range(self.n_epochs): + self.logger.info("Epoch%d:", step) + self.logger.info("training...") + self.train_epoch(train_loader) + self.logger.info("evaluating...") + train_loss, train_score = self.test_epoch(train_loader) + val_loss, val_score = self.test_epoch(valid_loader) + self.logger.info("train %.6f, valid %.6f" % (train_score, val_score)) + evals_result["train"].append(train_score) + evals_result["valid"].append(val_score) + + if val_score > best_score: + best_score = val_score + stop_steps = 0 + best_epoch = step + best_param = copy.deepcopy(self.ALSTM_model.state_dict()) + else: + stop_steps += 1 + if stop_steps >= self.early_stop: + self.logger.info("early stop") + break + + self.logger.info("best score: %.6lf @ %d" % (best_score, best_epoch)) + self.ALSTM_model.load_state_dict(best_param) + torch.save(best_param, save_path) + + if self.use_gpu: + torch.cuda.empty_cache() + + def predict(self, dataset): + if not self._fitted: + raise ValueError("model is not fitted yet!") + + dl_test = dataset.prepare("test", data_key=DataHandlerLP.DK_I) + dl_test.config(fillna_type="ffill+bfill") + test_loader = DataLoader(dl_test, batch_size=self.batch_size, num_workers=20) + self.ALSTM_model.eval() + preds = [] + + for data in test_loader: + + feature = data[:,:,0:-1].to(self.device) + + with torch.no_grad(): + if self.use_gpu: + pred = self.ALSTM_model(feature.float()).detach().cpu().numpy() + else: + pred = self.ALSTM_model(feature.float()).detach().numpy() + + preds.append(pred) + + return pd.Series(np.concatenate(preds), index=dl_test.get_index()) + + +class ALSTMModel(nn.Module): + def __init__(self, d_feat=6, hidden_size=64, num_layers=2, dropout=0.0, rnn_type="GRU"): + super().__init__() + self.hid_size = hidden_size + self.input_size = d_feat + self.dropout = dropout + self.rnn_type = rnn_type + self.rnn_layer = num_layers + self._build_model() + + def _build_model(self): + try: + klass = getattr(nn, self.rnn_type.upper()) + except: + raise ValueError("unknown rnn_type `%s`" % self.rnn_type) + self.net = nn.Sequential() + self.net.add_module("fc_in", nn.Linear(in_features=self.input_size, out_features=self.hid_size)) + self.net.add_module("act", nn.Tanh()) + self.rnn = klass( + input_size=self.hid_size, + hidden_size=self.hid_size, + num_layers=self.rnn_layer, + batch_first=True, + dropout=self.dropout, + ) + self.fc_out = nn.Linear(in_features=self.hid_size * 2, out_features=1) + self.att_net = nn.Sequential() + self.att_net.add_module( + "att_fc_in", + nn.Linear(in_features=self.hid_size, out_features=int(self.hid_size / 2)), + ) + self.att_net.add_module("att_dropout", torch.nn.Dropout(self.dropout)) + self.att_net.add_module("att_act", nn.Tanh()) + self.att_net.add_module( + "att_fc_out", + nn.Linear(in_features=int(self.hid_size / 2), out_features=1, bias=False), + ) + self.att_net.add_module("att_softmax", nn.Softmax(dim=1)) + + def forward(self, inputs): + rnn_out, _ = self.rnn(self.net(inputs)) # [batch, seq_len, num_directions * hidden_size] + attention_score = self.att_net(rnn_out) # [batch, seq_len, 1] + out_att = torch.mul(rnn_out, attention_score) + out_att = torch.sum(out_att, dim=1) + out = self.fc_out( + torch.cat((rnn_out[:, -1, :], out_att), dim=1) + ) # [batch, seq_len, num_directions * hidden_size] -> [batch, 1] + return out[..., 0] diff --git a/qlib/contrib/model/pytorch_gats_ts.py b/qlib/contrib/model/pytorch_gats_ts.py new file mode 100644 index 000000000..a2ff154de --- /dev/null +++ b/qlib/contrib/model/pytorch_gats_ts.py @@ -0,0 +1,379 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + + +from __future__ import division +from __future__ import print_function + +import os +import numpy as np +import pandas as pd +import copy +from sklearn.metrics import roc_auc_score, mean_squared_error +import logging +from ...utils import ( + unpack_archive_with_buffer, + save_multiple_parts_file, + create_save_path, + drop_nan_by_y_index, +) +from ...log import get_module_logger, TimeInspector +import torch +import torch.nn as nn +import torch.optim as optim + +from ...model.base import Model +from ...data.dataset import DatasetH +from ...data.dataset.handler import DataHandlerLP +from ...contrib.model.pytorch_lstm import LSTMModel +from ...contrib.model.pytorch_gru import GRUModel + + +class GATs(Model): + """GATs Model + + Parameters + ---------- + lr : float + learning rate + d_feat : int + input dimensions for each time step + metric : str + the evaluate metric used in early stop + optimizer : str + optimizer name + GPU : str + the GPU ID(s) used for training + """ + + def __init__( + self, + d_feat=6, + hidden_size=64, + num_layers=2, + dropout=0.0, + n_epochs=200, + lr=0.001, + metric="", + early_stop=20, + loss="mse", + base_model="GRU", + with_pretrain=True, + model_path=None, + optimizer="adam", + GPU="0", + seed=None, + **kwargs + ): + # Set logger. + self.logger = get_module_logger("GATs") + self.logger.info("GATs pytorch version...") + + # set hyper-parameters. + self.d_feat = d_feat + self.hidden_size = hidden_size + self.num_layers = num_layers + self.dropout = dropout + self.n_epochs = n_epochs + self.lr = lr + self.metric = metric + self.early_stop = early_stop + self.optimizer = optimizer.lower() + self.loss = loss + self.base_model = base_model + self.with_pretrain = with_pretrain + self.model_path = model_path + self.device = torch.device("cuda:%d" % (GPU) if torch.cuda.is_available() else "cpu") + self.use_gpu = torch.cuda.is_available() + self.seed = seed + + self.logger.info( + "GATs parameters setting:" + "\nd_feat : {}" + "\nhidden_size : {}" + "\nnum_layers : {}" + "\ndropout : {}" + "\nn_epochs : {}" + "\nlr : {}" + "\nmetric : {}" + "\nearly_stop : {}" + "\noptimizer : {}" + "\nloss_type : {}" + "\nbase_model : {}" + "\nwith_pretrain : {}" + "\nmodel_path : {}" + "\nvisible_GPU : {}" + "\nuse_GPU : {}" + "\nseed : {}".format( + d_feat, + hidden_size, + num_layers, + dropout, + n_epochs, + lr, + metric, + early_stop, + optimizer.lower(), + loss, + base_model, + with_pretrain, + model_path, + GPU, + self.use_gpu, + seed, + ) + ) + + if self.seed is not None: + np.random.seed(self.seed) + torch.manual_seed(self.seed) + + self.GAT_model = GATModel( + d_feat=self.d_feat, + hidden_size=self.hidden_size, + num_layers=self.num_layers, + dropout=self.dropout, + base_model=self.base_model, + ) + if optimizer.lower() == "adam": + self.train_optimizer = optim.Adam(self.GAT_model.parameters(), lr=self.lr) + elif optimizer.lower() == "gd": + self.train_optimizer = optim.SGD(self.GAT_model.parameters(), lr=self.lr) + else: + raise NotImplementedError("optimizer {} is not supported!".format(optimizer)) + + self._fitted = False + self.GAT_model.to(self.device) + + def mse(self, pred, label): + loss = (pred - label) ** 2 + return torch.mean(loss) + + def loss_fn(self, pred, label): + mask = ~torch.isnan(label) + + if self.loss == "mse": + return self.mse(pred[mask], label[mask]) + + raise ValueError("unknown loss `%s`" % self.loss) + + def metric_fn(self, pred, label): + + mask = torch.isfinite(label) + + if self.metric == "" or self.metric == "loss": + return -self.loss_fn(pred[mask], label[mask]) + + raise ValueError("unknown metric `%s`" % self.metric) + + def get_daily_inter(self, df, shuffle=False): + # organize the train data into daily batches + daily_count = df.groupby(level=0).size().values + daily_index = np.roll(np.cumsum(daily_count), 1) + daily_index[0] = 0 + if shuffle: + # shuffle data + daily_shuffle = list(zip(daily_index, daily_count)) + np.random.shuffle(daily_shuffle) + daily_index, daily_count = zip(*daily_shuffle) + return daily_index, daily_count + + def train_epoch(self, data_loader): + + self.ALSTM_model.train() + + for data in data_loader: + feature = data[:,:,0:-1].to(self.device) + label = data[:,-1,-1].to(self.device) + + pred = self.ALSTM_model(feature.float()) + loss = self.loss_fn(pred, label) + + self.train_optimizer.zero_grad() + loss.backward() + torch.nn.utils.clip_grad_value_(self.ALSTM_model.parameters(), 3.0) + self.train_optimizer.step() + + def test_epoch(self, data_loader): + + self.ALSTM_model.eval() + + scores = [] + losses = [] + + for data in data_loader: + + feature = data[:,:,0:-1].to(self.device) + # feature[torch.isnan(feature)] = 0 + label = data[:,-1,-1].to(self.device) + + pred = self.ALSTM_model(feature.float()) + loss = self.loss_fn(pred, label) + losses.append(loss.item()) + + score = self.metric_fn(pred, label) + scores.append(score.item()) + + return np.mean(losses), np.mean(scores) + + def fit( + self, + dataset: DatasetH, + evals_result=dict(), + verbose=True, + save_path=None, + ): + + dl_train = dataset.prepare("train", data_key=DataHandlerLP.DK_L) + dl_valid = dataset.prepare("valid", data_key=DataHandlerLP.DK_L) + + dl_train.config(fillna_type="ffill+bfill") # process nan brought by dataloader + dl_valid.config(fillna_type="ffill+bfill") # process nan brought by dataloader + + train_loader = DataLoader(dl_train, batch_size=self.batch_size, shuffle=True, num_workers=self.n_jobs) + valid_loader = DataLoader(dl_valid, batch_size=self.batch_size, shuffle=False, num_workers=self.n_jobs) + + if save_path == None: + save_path = create_save_path(save_path) + + stop_steps = 0 + train_loss = 0 + best_score = -np.inf + best_epoch = 0 + evals_result["train"] = [] + evals_result["valid"] = [] + + # load pretrained base_model + if self.with_pretrain: + if self.model_path == None: + raise ValueError("the path of the pretrained model should be given first!") + self.logger.info("Loading pretrained model...") + if self.base_model == "LSTM": + pretrained_model = LSTMModel() + pretrained_model.load_state_dict(torch.load(self.model_path)) + elif self.base_model == "GRU": + pretrained_model = GRUModel() + pretrained_model.load_state_dict(torch.load(self.model_path)) + else: + raise ValueError("unknown base model name `%s`" % self.base_model) + + model_dict = self.GAT_model.state_dict() + pretrained_dict = {k: v for k, v in pretrained_model.state_dict().items() if k in model_dict} + model_dict.update(pretrained_dict) + self.GAT_model.load_state_dict(model_dict) + self.logger.info("Loading pretrained model Done...") + + # train + self.logger.info("training...") + self._fitted = True + + for step in range(self.n_epochs): + self.logger.info("Epoch%d:", step) + self.logger.info("training...") + self.train_epoch(train_loader) + self.logger.info("evaluating...") + train_loss, train_score = self.test_epoch(train_loader) + val_loss, val_score = self.test_epoch(valid_loader) + self.logger.info("train %.6f, valid %.6f" % (train_score, val_score)) + evals_result["train"].append(train_score) + evals_result["valid"].append(val_score) + + if val_score > best_score: + best_score = val_score + stop_steps = 0 + best_epoch = step + best_param = copy.deepcopy(self.ALSTM_model.state_dict()) + else: + stop_steps += 1 + if stop_steps >= self.early_stop: + self.logger.info("early stop") + break + + self.logger.info("best score: %.6lf @ %d" % (best_score, best_epoch)) + self.ALSTM_model.load_state_dict(best_param) + torch.save(best_param, save_path) + + if self.use_gpu: + torch.cuda.empty_cache() + + def predict(self, dataset): + if not self._fitted: + raise ValueError("model is not fitted yet!") + + dl_test = dataset.prepare("test", data_key=DataHandlerLP.DK_I) + dl_test.config(fillna_type="ffill+bfill") + test_loader = DataLoader(dl_test, batch_size=self.batch_size, num_workers=20) + self.ALSTM_model.eval() + preds = [] + + for data in test_loader: + + feature = data[:,:,0:-1].to(self.device) + + with torch.no_grad(): + if self.use_gpu: + pred = self.ALSTM_model(feature.float()).detach().cpu().numpy() + else: + pred = self.ALSTM_model(feature.float()).detach().numpy() + + preds.append(pred) + + return pd.Series(np.concatenate(preds), index=dl_test.get_index()) + + +class GATModel(nn.Module): + def __init__(self, d_feat=6, hidden_size=64, num_layers=2, dropout=0.0, base_model="GRU"): + super().__init__() + + if base_model == "GRU": + self.rnn = nn.GRU( + input_size=d_feat, + hidden_size=hidden_size, + num_layers=num_layers, + batch_first=True, + dropout=dropout, + ) + elif base_model == "LSTM": + self.rnn = nn.LSTM( + input_size=d_feat, + hidden_size=hidden_size, + num_layers=num_layers, + batch_first=True, + dropout=dropout, + ) + else: + raise ValueError("unknown base model name `%s`" % base_model) + + self.hidden_size = hidden_size + self.d_feat = d_feat + self.transformation = nn.Linear(self.hidden_size, self.hidden_size) + self.a = nn.Parameter(torch.randn(self.hidden_size * 2, 1)) + self.a.requires_grad = True + self.fc = nn.Linear(self.hidden_size, self.hidden_size) + self.fc_out = nn.Linear(hidden_size, 1) + self.leaky_relu = nn.LeakyReLU() + self.softmax = nn.Softmax(dim=1) + + def cal_attention(self, x, y): + x = self.transformation(x) + y = self.transformation(y) + + sample_num = x.shape[0] + dim = x.shape[1] + e_x = x.expand(sample_num, sample_num, dim) + e_y = torch.transpose(e_x, 0, 1) + attention_in = torch.cat((e_x, e_y), 2).view(-1, dim * 2) + self.a_t = torch.t(self.a) + attention_out = self.a_t.mm(torch.t(attention_in)).view(sample_num, sample_num) + attention_out = self.leaky_relu(attention_out) + att_weight = self.softmax(attention_out) + return att_weight + + def forward(self, x): + out, _ = self.rnn(x) + hidden = out[:, -1, :] + att_weight = self.cal_attention(hidden, hidden) + hidden = att_weight.mm(hidden) + hidden + hidden = self.fc(hidden) + hidden = self.leaky_relu(hidden) + return self.fc_out(hidden).squeeze() diff --git a/qlib/contrib/model/pytorch_gru_ts.py b/qlib/contrib/model/pytorch_gru_ts.py new file mode 100755 index 000000000..217aea0ae --- /dev/null +++ b/qlib/contrib/model/pytorch_gru_ts.py @@ -0,0 +1,302 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + + +from __future__ import division +from __future__ import print_function + +import os +import numpy as np +import pandas as pd +import copy +from sklearn.metrics import roc_auc_score, mean_squared_error +import logging +from ...utils import ( + unpack_archive_with_buffer, + save_multiple_parts_file, + create_save_path, + drop_nan_by_y_index, +) +from ...log import get_module_logger, TimeInspector + +import torch +import torch.nn as nn +import torch.optim as optim +from torch.utils.data import DataLoader + +from ...model.base import Model +from ...data.dataset import DatasetH, TSDatasetH +from ...data.dataset.handler import DataHandlerLP + + +class GRU(Model): + """GRU Model + + Parameters + ---------- + d_feat : int + input dimension for each time step + metric: str + the evaluate metric used in early stop + optimizer : str + optimizer name + GPU : str + the GPU ID(s) used for training + """ + + def __init__( + self, + d_feat=6, + hidden_size=64, + num_layers=2, + dropout=0.0, + n_epochs=200, + lr=0.001, + metric="", + batch_size=2000, + early_stop=20, + loss="mse", + optimizer="adam", + n_jobs=10, + GPU="0", + seed=None, + **kwargs + ): + # Set logger. + self.logger = get_module_logger("GRU") + self.logger.info("GRU pytorch version...") + + # set hyper-parameters. + self.d_feat = d_feat + self.hidden_size = hidden_size + self.num_layers = num_layers + self.dropout = dropout + self.n_epochs = n_epochs + self.lr = lr + self.metric = metric + self.batch_size = batch_size + self.early_stop = early_stop + self.optimizer = optimizer.lower() + self.loss = loss + self.device = torch.device("cuda:%d" % (GPU) if torch.cuda.is_available() else "cpu") + self.n_jobs = n_jobs + self.use_gpu = torch.cuda.is_available() + self.seed = seed + + self.logger.info( + "GRU parameters setting:" + "\nd_feat : {}" + "\nhidden_size : {}" + "\nnum_layers : {}" + "\ndropout : {}" + "\nn_epochs : {}" + "\nlr : {}" + "\nmetric : {}" + "\nbatch_size : {}" + "\nearly_stop : {}" + "\noptimizer : {}" + "\nloss_type : {}" + "\nvisible_GPU : {}" + "\nn_jobs : {}" + "\nuse_GPU : {}" + "\nseed : {}".format( + d_feat, + hidden_size, + num_layers, + dropout, + n_epochs, + lr, + metric, + batch_size, + early_stop, + optimizer.lower(), + loss, + GPU, + n_jobs, + self.use_gpu, + seed, + ) + ) + + if self.seed is not None: + np.random.seed(self.seed) + torch.manual_seed(self.seed) + + self.GRU_model = GRUModel( + d_feat=self.d_feat, + hidden_size=self.hidden_size, + num_layers=self.num_layers, + dropout=self.dropout, + ).to(self.device) + if optimizer.lower() == "adam": + self.train_optimizer = optim.Adam(self.GRU_model.parameters(), lr=self.lr) + elif optimizer.lower() == "gd": + self.train_optimizer = optim.SGD(self.GRU_model.parameters(), lr=self.lr) + else: + raise NotImplementedError("optimizer {} is not supported!".format(optimizer)) + + self._fitted = False + self.GRU_model.to(self.device) + + def mse(self, pred, label): + loss = (pred - label) ** 2 + return torch.mean(loss) + + def loss_fn(self, pred, label): + mask = ~torch.isnan(label) + + if self.loss == "mse": + return self.mse(pred[mask], label[mask]) + + raise ValueError("unknown loss `%s`" % self.loss) + + def metric_fn(self, pred, label): + + mask = torch.isfinite(label) + + if self.metric == "" or self.metric == "loss": + return -self.loss_fn(pred[mask], label[mask]) + + raise ValueError("unknown metric `%s`" % self.metric) + + def train_epoch(self, data_loader): + + self.GRU_model.train() + + for data in data_loader: + feature = data[:,:,0:-1].to(self.device) + label = data[:,-1,-1].to(self.device) + + pred = self.GRU_model(feature.float()) + loss = self.loss_fn(pred, label) + + self.train_optimizer.zero_grad() + loss.backward() + torch.nn.utils.clip_grad_value_(self.GRU_model.parameters(), 3.0) + self.train_optimizer.step() + + def test_epoch(self, data_loader): + + self.GRU_model.eval() + + scores = [] + losses = [] + + for data in data_loader: + + feature = data[:,:,0:-1].to(self.device) + # feature[torch.isnan(feature)] = 0 + label = data[:,-1,-1].to(self.device) + + pred = self.GRU_model(feature.float()) + loss = self.loss_fn(pred, label) + losses.append(loss.item()) + + score = self.metric_fn(pred, label) + scores.append(score.item()) + + return np.mean(losses), np.mean(scores) + + def fit( + self, + dataset, + evals_result=dict(), + verbose=True, + save_path=None, + ): + dl_train = dataset.prepare("train", data_key=DataHandlerLP.DK_L) + dl_valid = dataset.prepare("valid", data_key=DataHandlerLP.DK_L) + + dl_train.config(fillna_type="ffill+bfill") # process nan brought by dataloader + dl_valid.config(fillna_type="ffill+bfill") # process nan brought by dataloader + + train_loader = DataLoader(dl_train, batch_size=self.batch_size, shuffle=True, num_workers=self.n_jobs) + valid_loader = DataLoader(dl_valid, batch_size=self.batch_size, shuffle=False, num_workers=self.n_jobs) + + if save_path == None: + save_path = create_save_path(save_path) + + stop_steps = 0 + train_loss = 0 + best_score = -np.inf + best_epoch = 0 + evals_result["train"] = [] + evals_result["valid"] = [] + + # train + self.logger.info("training...") + self._fitted = True + + for step in range(self.n_epochs): + self.logger.info("Epoch%d:", step) + self.logger.info("training...") + self.train_epoch(train_loader) + self.logger.info("evaluating...") + train_loss, train_score = self.test_epoch(train_loader) + val_loss, val_score = self.test_epoch(valid_loader) + self.logger.info("train %.6f, valid %.6f" % (train_score, val_score)) + evals_result["train"].append(train_score) + evals_result["valid"].append(val_score) + + if val_score > best_score: + best_score = val_score + stop_steps = 0 + best_epoch = step + best_param = copy.deepcopy(self.GRU_model.state_dict()) + else: + stop_steps += 1 + if stop_steps >= self.early_stop: + self.logger.info("early stop") + break + + self.logger.info("best score: %.6lf @ %d" % (best_score, best_epoch)) + self.GRU_model.load_state_dict(best_param) + # torch.save(best_param, save_path) + torch.save(best_param, '/home/lewwang/qlib/examples/benchmarks/GRU/csi300_gru_ts.pkl') + + if self.use_gpu: + torch.cuda.empty_cache() + + def predict(self, dataset): + if not self._fitted: + raise ValueError("model is not fitted yet!") + + dl_test = dataset.prepare("test", data_key=DataHandlerLP.DK_I) + dl_test.config(fillna_type="ffill+bfill") + test_loader = DataLoader(dl_test, batch_size=self.batch_size, num_workers=20) + self.GRU_model.eval() + preds = [] + + for data in test_loader: + + feature = data[:,:,0:-1].to(self.device) + + with torch.no_grad(): + if self.use_gpu: + pred = self.GRU_model(feature.float()).detach().cpu().numpy() + else: + pred = self.GRU_model(feature.float()).detach().numpy() + + preds.append(pred) + + return pd.Series(np.concatenate(preds), index=dl_test.get_index()) + + +class GRUModel(nn.Module): + def __init__(self, d_feat=6, hidden_size=64, num_layers=2, dropout=0.0): + super().__init__() + + self.rnn = nn.GRU( + input_size=d_feat, + hidden_size=hidden_size, + num_layers=num_layers, + batch_first=True, + dropout=dropout, + ) + self.fc_out = nn.Linear(hidden_size, 1) + + self.d_feat = d_feat + + def forward(self, x): + out, _ = self.rnn(x) + return self.fc_out(out[:, -1, :]).squeeze() diff --git a/qlib/contrib/model/pytorch_lstm_ts.py b/qlib/contrib/model/pytorch_lstm_ts.py new file mode 100755 index 000000000..676850fb2 --- /dev/null +++ b/qlib/contrib/model/pytorch_lstm_ts.py @@ -0,0 +1,303 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + + +from __future__ import division +from __future__ import print_function + +import os +import numpy as np +import pandas as pd +import copy +from sklearn.metrics import roc_auc_score, mean_squared_error +import logging +from ...utils import ( + unpack_archive_with_buffer, + save_multiple_parts_file, + create_save_path, + drop_nan_by_y_index, +) +from ...log import get_module_logger, TimeInspector + +import torch +import torch.nn as nn +import torch.optim as optim +from torch.utils.data import DataLoader + +from ...model.base import Model +from ...data.dataset import DatasetH, TSDatasetH +from ...data.dataset.handler import DataHandlerLP + + +class LSTM(Model): + """LSTM Model + + Parameters + ---------- + d_feat : int + input dimension for each time step + metric: str + the evaluate metric used in early stop + optimizer : str + optimizer name + GPU : str + the GPU ID(s) used for training + """ + + def __init__( + self, + d_feat=6, + hidden_size=64, + num_layers=2, + dropout=0.0, + n_epochs=200, + lr=0.001, + metric="", + batch_size=2000, + early_stop=20, + loss="mse", + optimizer="adam", + n_jobs=10, + GPU="0", + seed=None, + **kwargs + ): + # Set logger. + self.logger = get_module_logger("LSTM") + self.logger.info("LSTM pytorch version...") + + # set hyper-parameters. + self.d_feat = d_feat + self.hidden_size = hidden_size + self.num_layers = num_layers + self.dropout = dropout + self.n_epochs = n_epochs + self.lr = lr + self.metric = metric + self.batch_size = batch_size + self.early_stop = early_stop + self.optimizer = optimizer.lower() + self.loss = loss + self.device = torch.device("cuda:%d" % (GPU) if torch.cuda.is_available() else "cpu") + self.n_jobs = n_jobs + self.use_gpu = torch.cuda.is_available() + self.seed = seed + + self.logger.info( + "LSTM parameters setting:" + "\nd_feat : {}" + "\nhidden_size : {}" + "\nnum_layers : {}" + "\ndropout : {}" + "\nn_epochs : {}" + "\nlr : {}" + "\nmetric : {}" + "\nbatch_size : {}" + "\nearly_stop : {}" + "\noptimizer : {}" + "\nloss_type : {}" + "\nvisible_GPU : {}" + "\nn_jobs : {}" + "\nuse_GPU : {}" + "\nseed : {}".format( + d_feat, + hidden_size, + num_layers, + dropout, + n_epochs, + lr, + metric, + batch_size, + early_stop, + optimizer.lower(), + loss, + GPU, + n_jobs, + self.use_gpu, + seed, + ) + ) + + if self.seed is not None: + np.random.seed(self.seed) + torch.manual_seed(self.seed) + + self.LSTM_model = LSTMModel( + d_feat=self.d_feat, + hidden_size=self.hidden_size, + num_layers=self.num_layers, + dropout=self.dropout, + ).to(self.device) + if optimizer.lower() == "adam": + self.train_optimizer = optim.Adam(self.LSTM_model.parameters(), lr=self.lr) + elif optimizer.lower() == "gd": + self.train_optimizer = optim.SGD(self.LSTM_model.parameters(), lr=self.lr) + else: + raise NotImplementedError("optimizer {} is not supported!".format(optimizer)) + + self._fitted = False + self.LSTM_model.to(self.device) + + def mse(self, pred, label): + loss = (pred - label) ** 2 + return torch.mean(loss) + + def loss_fn(self, pred, label): + mask = ~torch.isnan(label) + + if self.loss == "mse": + return self.mse(pred[mask], label[mask]) + + raise ValueError("unknown loss `%s`" % self.loss) + + def metric_fn(self, pred, label): + + mask = torch.isfinite(label) + + if self.metric == "" or self.metric == "loss": + return -self.loss_fn(pred[mask], label[mask]) + + raise ValueError("unknown metric `%s`" % self.metric) + + def train_epoch(self, data_loader): + + self.LSTM_model.train() + + for data in data_loader: + feature = data[:,:,0:-1].to(self.device) + label = data[:,-1,-1].to(self.device) + + pred = self.LSTM_model(feature.float()) + loss = self.loss_fn(pred, label) + + self.train_optimizer.zero_grad() + loss.backward() + torch.nn.utils.clip_grad_value_(self.LSTM_model.parameters(), 3.0) + self.train_optimizer.step() + + def test_epoch(self, data_loader): + + self.LSTM_model.eval() + + scores = [] + losses = [] + + for data in data_loader: + + feature = data[:,:,0:-1].to(self.device) + # feature[torch.isnan(feature)] = 0 + label = data[:,-1,-1].to(self.device) + + pred = self.LSTM_model(feature.float()) + loss = self.loss_fn(pred, label) + losses.append(loss.item()) + + score = self.metric_fn(pred, label) + scores.append(score.item()) + + return np.mean(losses), np.mean(scores) + + def fit( + self, + dataset, + evals_result=dict(), + verbose=True, + save_path=None, + ): + dl_train = dataset.prepare("train", data_key=DataHandlerLP.DK_L) + dl_valid = dataset.prepare("valid", data_key=DataHandlerLP.DK_L) + + dl_train.config(fillna_type="ffill+bfill") # process nan brought by dataloader + dl_valid.config(fillna_type="ffill+bfill") # process nan brought by dataloader + + train_loader = DataLoader(dl_train, batch_size=self.batch_size, shuffle=True, num_workers=self.n_jobs) + valid_loader = DataLoader(dl_valid, batch_size=self.batch_size, shuffle=False, num_workers=self.n_jobs) + + if save_path == None: + save_path = create_save_path(save_path) + + stop_steps = 0 + train_loss = 0 + best_score = -np.inf + best_epoch = 0 + evals_result["train"] = [] + evals_result["valid"] = [] + + # train + self.logger.info("training...") + self._fitted = True + + for step in range(self.n_epochs): + self.logger.info("Epoch%d:", step) + self.logger.info("training...") + self.train_epoch(train_loader) + self.logger.info("evaluating...") + train_loss, train_score = self.test_epoch(train_loader) + val_loss, val_score = self.test_epoch(valid_loader) + self.logger.info("train %.6f, valid %.6f" % (train_score, val_score)) + evals_result["train"].append(train_score) + evals_result["valid"].append(val_score) + + if val_score > best_score: + best_score = val_score + stop_steps = 0 + best_epoch = step + best_param = copy.deepcopy(self.LSTM_model.state_dict()) + else: + stop_steps += 1 + if stop_steps >= self.early_stop: + self.logger.info("early stop") + break + + self.logger.info("best score: %.6lf @ %d" % (best_score, best_epoch)) + self.LSTM_model.load_state_dict(best_param) + # torch.save(best_param, save_path) + torch.save(best_param, '/home/lewwang/qlib/examples/benchmarks/LSTM/csi300_lstm_ts.pkl') + + + if self.use_gpu: + torch.cuda.empty_cache() + + def predict(self, dataset): + if not self._fitted: + raise ValueError("model is not fitted yet!") + + dl_test = dataset.prepare("test", data_key=DataHandlerLP.DK_I) + dl_test.config(fillna_type="ffill+bfill") + test_loader = DataLoader(dl_test, batch_size=self.batch_size, num_workers=20) + self.LSTM_model.eval() + preds = [] + + for data in test_loader: + + feature = data[:,:,0:-1].to(self.device) + + with torch.no_grad(): + if self.use_gpu: + pred = self.LSTM_model(feature.float()).detach().cpu().numpy() + else: + pred = self.LSTM_model(feature.float()).detach().numpy() + + preds.append(pred) + + return pd.Series(np.concatenate(preds), index=dl_test.get_index()) + + +class LSTMModel(nn.Module): + def __init__(self, d_feat=6, hidden_size=64, num_layers=2, dropout=0.0): + super().__init__() + + self.rnn = nn.LSTM( + input_size=d_feat, + hidden_size=hidden_size, + num_layers=num_layers, + batch_first=True, + dropout=dropout, + ) + self.fc_out = nn.Linear(hidden_size, 1) + + self.d_feat = d_feat + + def forward(self, x): + out, _ = self.rnn(x) + return self.fc_out(out[:, -1, :]).squeeze() \ No newline at end of file diff --git a/qlib/contrib/model/pytorch_sfm_ts.py b/qlib/contrib/model/pytorch_sfm_ts.py new file mode 100644 index 000000000..fcf5a0f75 --- /dev/null +++ b/qlib/contrib/model/pytorch_sfm_ts.py @@ -0,0 +1,434 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + + +from __future__ import division +from __future__ import print_function + +import os +import numpy as np +import pandas as pd +import copy +from sklearn.metrics import roc_auc_score, mean_squared_error +import logging +from ...utils import ( + unpack_archive_with_buffer, + save_multiple_parts_file, + create_save_path, + drop_nan_by_y_index, +) +from ...log import get_module_logger, TimeInspector + +import torch +import torch.nn as nn +import torch.optim as optim +from torch.utils.data import DataLoader + +from ...model.base import Model +from ...data.dataset import DatasetH, TSDatasetH +from ...data.dataset.handler import DataHandlerLP + + +class SFM(Model): + """SFM Model + + Parameters + ---------- + d_feat : int + input dimension for each time step + metric: str + the evaluate metric used in early stop + optimizer : str + optimizer name + GPU : str + the GPU ID(s) used for training + """ + + def __init__( + self, + d_feat=6, + hidden_size=64, + num_layers=2, + dropout=0.0, + n_epochs=200, + lr=0.001, + metric="", + batch_size=2000, + early_stop=20, + loss="mse", + optimizer="adam", + n_jobs=10, + GPU="0", + seed=None, + **kwargs + ): + # Set logger. + self.logger = get_module_logger("SFM") + self.logger.info("SFM pytorch version...") + + # set hyper-parameters. + self.d_feat = d_feat + self.hidden_size = hidden_size + self.num_layers = num_layers + self.dropout = dropout + self.n_epochs = n_epochs + self.lr = lr + self.metric = metric + self.batch_size = batch_size + self.early_stop = early_stop + self.optimizer = optimizer.lower() + self.loss = loss + self.device = torch.device("cuda:%d" % (GPU) if torch.cuda.is_available() else "cpu") + self.n_jobs = n_jobs + self.use_gpu = torch.cuda.is_available() + self.seed = seed + + self.logger.info( + "SFM parameters setting:" + "\nd_feat : {}" + "\nhidden_size : {}" + "\nnum_layers : {}" + "\ndropout : {}" + "\nn_epochs : {}" + "\nlr : {}" + "\nmetric : {}" + "\nbatch_size : {}" + "\nearly_stop : {}" + "\noptimizer : {}" + "\nloss_type : {}" + "\nvisible_GPU : {}" + "\nn_jobs : {}" + "\nuse_GPU : {}" + "\nseed : {}".format( + d_feat, + hidden_size, + num_layers, + dropout, + n_epochs, + lr, + metric, + batch_size, + early_stop, + optimizer.lower(), + loss, + GPU, + n_jobs, + self.use_gpu, + seed, + ) + ) + + if self.seed is not None: + np.random.seed(self.seed) + torch.manual_seed(self.seed) + + self.SFM_model = SFMModel( + d_feat=self.d_feat, + hidden_size=self.hidden_size, + num_layers=self.num_layers, + dropout=self.dropout, + ).to(self.device) + if optimizer.lower() == "adam": + self.train_optimizer = optim.Adam(self.SFM_model.parameters(), lr=self.lr) + elif optimizer.lower() == "gd": + self.train_optimizer = optim.SGD(self.SFM_model.parameters(), lr=self.lr) + else: + raise NotImplementedError("optimizer {} is not supported!".format(optimizer)) + + self._fitted = False + self.SFM_model.to(self.device) + + def mse(self, pred, label): + loss = (pred - label) ** 2 + return torch.mean(loss) + + def loss_fn(self, pred, label): + mask = ~torch.isnan(label) + + if self.loss == "mse": + return self.mse(pred[mask], label[mask]) + + raise ValueError("unknown loss `%s`" % self.loss) + + def metric_fn(self, pred, label): + + mask = torch.isfinite(label) + + if self.metric == "" or self.metric == "loss": + return -self.loss_fn(pred[mask], label[mask]) + + raise ValueError("unknown metric `%s`" % self.metric) + + def train_epoch(self, data_loader): + + self.SFM_model.train() + + for data in data_loader: + feature = data[:,:,0:-1].to(self.device) + label = data[:,-1,-1].to(self.device) + + pred = self.SFM_model(feature.float()) + loss = self.loss_fn(pred, label) + + self.train_optimizer.zero_grad() + loss.backward() + torch.nn.utils.clip_grad_value_(self.SFM_model.parameters(), 3.0) + self.train_optimizer.step() + + def test_epoch(self, data_loader): + + self.SFM_model.eval() + + scores = [] + losses = [] + + for data in data_loader: + + feature = data[:,:,0:-1].to(self.device) + # feature[torch.isnan(feature)] = 0 + label = data[:,-1,-1].to(self.device) + + pred = self.SFM_model(feature.float()) + loss = self.loss_fn(pred, label) + losses.append(loss.item()) + + score = self.metric_fn(pred, label) + scores.append(score.item()) + + return np.mean(losses), np.mean(scores) + + def fit( + self, + dataset, + evals_result=dict(), + verbose=True, + save_path=None, + ): + dl_train = dataset.prepare("train", data_key=DataHandlerLP.DK_L) + dl_valid = dataset.prepare("valid", data_key=DataHandlerLP.DK_L) + + dl_train.config(fillna_type="ffill+bfill") # process nan brought by dataloader + dl_valid.config(fillna_type="ffill+bfill") # process nan brought by dataloader + + train_loader = DataLoader(dl_train, batch_size=self.batch_size, shuffle=True, num_workers=self.n_jobs) + valid_loader = DataLoader(dl_valid, batch_size=self.batch_size, shuffle=False, num_workers=self.n_jobs) + + if save_path == None: + save_path = create_save_path(save_path) + + stop_steps = 0 + train_loss = 0 + best_score = -np.inf + best_epoch = 0 + evals_result["train"] = [] + evals_result["valid"] = [] + + # train + self.logger.info("training...") + self._fitted = True + + for step in range(self.n_epochs): + self.logger.info("Epoch%d:", step) + self.logger.info("training...") + self.train_epoch(train_loader) + self.logger.info("evaluating...") + train_loss, train_score = self.test_epoch(train_loader) + val_loss, val_score = self.test_epoch(valid_loader) + self.logger.info("train %.6f, valid %.6f" % (train_score, val_score)) + evals_result["train"].append(train_score) + evals_result["valid"].append(val_score) + + if val_score > best_score: + best_score = val_score + stop_steps = 0 + best_epoch = step + best_param = copy.deepcopy(self.SFM_model.state_dict()) + else: + stop_steps += 1 + if stop_steps >= self.early_stop: + self.logger.info("early stop") + break + + self.logger.info("best score: %.6lf @ %d" % (best_score, best_epoch)) + self.SFM_model.load_state_dict(best_param) + torch.save(best_param, save_path) + + if self.use_gpu: + torch.cuda.empty_cache() + + def predict(self, dataset): + if not self._fitted: + raise ValueError("model is not fitted yet!") + + dl_test = dataset.prepare("test", data_key=DataHandlerLP.DK_I) + dl_test.config(fillna_type="ffill+bfill") + test_loader = DataLoader(dl_test, batch_size=self.batch_size, num_workers=20) + self.SFM_model.eval() + preds = [] + + for data in test_loader: + + feature = data[:,:,0:-1].to(self.device) + + with torch.no_grad(): + if self.use_gpu: + pred = self.SFM_model(feature.float()).detach().cpu().numpy() + else: + pred = self.SFM_model(feature.float()).detach().numpy() + + preds.append(pred) + + return pd.Series(np.concatenate(preds), index=dl_test.get_index()) + + +class SFM_Model(nn.Module): + def __init__( + self, + d_feat=6, + output_dim=1, + freq_dim=10, + hidden_size=64, + dropout_W=0.0, + dropout_U=0.0, + device="cpu", + ): + super().__init__() + + self.input_dim = d_feat + self.output_dim = output_dim + self.freq_dim = freq_dim + self.hidden_dim = hidden_size + self.device = device + + self.W_i = nn.Parameter(init.xavier_uniform_(torch.empty((self.input_dim, self.hidden_dim)))) + self.U_i = nn.Parameter(init.orthogonal_(torch.empty(self.hidden_dim, self.hidden_dim))) + self.b_i = nn.Parameter(torch.zeros(self.hidden_dim)) + + self.W_ste = nn.Parameter(init.xavier_uniform_(torch.empty(self.input_dim, self.hidden_dim))) + self.U_ste = nn.Parameter(init.orthogonal_(torch.empty(self.hidden_dim, self.hidden_dim))) + self.b_ste = nn.Parameter(torch.ones(self.hidden_dim)) + + self.W_fre = nn.Parameter(init.xavier_uniform_(torch.empty(self.input_dim, self.freq_dim))) + self.U_fre = nn.Parameter(init.orthogonal_(torch.empty(self.hidden_dim, self.freq_dim))) + self.b_fre = nn.Parameter(torch.ones(self.freq_dim)) + + self.W_c = nn.Parameter(init.xavier_uniform_(torch.empty(self.input_dim, self.hidden_dim))) + self.U_c = nn.Parameter(init.orthogonal_(torch.empty(self.hidden_dim, self.hidden_dim))) + self.b_c = nn.Parameter(torch.zeros(self.hidden_dim)) + + self.W_o = nn.Parameter(init.xavier_uniform_(torch.empty(self.input_dim, self.hidden_dim))) + self.U_o = nn.Parameter(init.orthogonal_(torch.empty(self.hidden_dim, self.hidden_dim))) + self.b_o = nn.Parameter(torch.zeros(self.hidden_dim)) + + self.U_a = nn.Parameter(init.orthogonal_(torch.empty(self.freq_dim, 1))) + self.b_a = nn.Parameter(torch.zeros(self.hidden_dim)) + + self.W_p = nn.Parameter(init.xavier_uniform_(torch.empty(self.hidden_dim, self.output_dim))) + self.b_p = nn.Parameter(torch.zeros(self.output_dim)) + + self.activation = nn.Tanh() + self.inner_activation = nn.Hardsigmoid() + self.dropout_W, self.dropout_U = (dropout_W, dropout_U) + self.fc_out = nn.Linear(self.output_dim, 1) + + self.states = [] + + def forward(self, input): + time_step = input.shape[1] + + for ts in range(time_step): + x = input[:, ts, :] + if len(self.states) == 0: # hasn't initialized yet + self.init_states(x) + self.get_constants(x) + p_tm1 = self.states[0] + h_tm1 = self.states[1] + S_re_tm1 = self.states[2] + S_im_tm1 = self.states[3] + time_tm1 = self.states[4] + B_U = self.states[5] + B_W = self.states[6] + frequency = self.states[7] + + x_i = torch.matmul(x * B_W[0], self.W_i) + self.b_i + x_ste = torch.matmul(x * B_W[0], self.W_ste) + self.b_ste + x_fre = torch.matmul(x * B_W[0], self.W_fre) + self.b_fre + x_c = torch.matmul(x * B_W[0], self.W_c) + self.b_c + x_o = torch.matmul(x * B_W[0], self.W_o) + self.b_o + + i = self.inner_activation(x_i + torch.matmul(h_tm1 * B_U[0], self.U_i)) + ste = self.inner_activation(x_ste + torch.matmul(h_tm1 * B_U[0], self.U_ste)) + fre = self.inner_activation(x_fre + torch.matmul(h_tm1 * B_U[0], self.U_fre)) + + ste = torch.reshape(ste, (-1, self.hidden_dim, 1)) + fre = torch.reshape(fre, (-1, 1, self.freq_dim)) + + f = ste * fre + + c = i * self.activation(x_c + torch.matmul(h_tm1 * B_U[0], self.U_c)) + + time = time_tm1 + 1 + + omega = torch.tensor(2 * np.pi) * time * frequency + + re = torch.cos(omega) + im = torch.sin(omega) + + c = torch.reshape(c, (-1, self.hidden_dim, 1)) + + S_re = f * S_re_tm1 + c * re + S_im = f * S_im_tm1 + c * im + + A = torch.square(S_re) + torch.square(S_im) + + A = torch.reshape(A, (-1, self.freq_dim)).float() + A_a = torch.matmul(A * B_U[0], self.U_a) + A_a = torch.reshape(A_a, (-1, self.hidden_dim)) + a = self.activation(A_a + self.b_a) + + o = self.inner_activation(x_o + torch.matmul(h_tm1 * B_U[0], self.U_o)) + + h = o * a + p = torch.matmul(h, self.W_p) + self.b_p + + self.states = [p, h, S_re, S_im, time, None, None, None] + self.states = [] + return self.fc_out(p).squeeze() + + def init_states(self, x): + reducer_f = torch.zeros((self.hidden_dim, self.freq_dim)).to(self.device) + reducer_p = torch.zeros((self.hidden_dim, self.output_dim)).to(self.device) + + init_state_h = torch.zeros(self.hidden_dim).to(self.device) + init_state_p = torch.matmul(init_state_h, reducer_p) + + init_state = torch.zeros_like(init_state_h).to(self.device) + init_freq = torch.matmul(init_state_h, reducer_f) + + init_state = torch.reshape(init_state, (-1, self.hidden_dim, 1)) + init_freq = torch.reshape(init_freq, (-1, 1, self.freq_dim)) + + init_state_S_re = init_state * init_freq + init_state_S_im = init_state * init_freq + + init_state_time = torch.tensor(0).to(self.device) + + self.states = [ + init_state_p, + init_state_h, + init_state_S_re, + init_state_S_im, + init_state_time, + None, + None, + None, + ] + + def get_constants(self, x): + constants = [] + constants.append([torch.tensor(1.0).to(self.device) for _ in range(6)]) + constants.append([torch.tensor(1.0).to(self.device) for _ in range(7)]) + array = np.array([float(ii) / self.freq_dim for ii in range(self.freq_dim)]) + constants.append(torch.tensor(array).to(self.device)) + + self.states[5:] = constants \ No newline at end of file