liuyu/qlib
1
0
Fork 0
mirror of https://github.com/microsoft/qlib.git synced 2026-07-17 01:14:35 +08:00
Code Activity

Update TCTS. (#495)

Browse Source 
* Update TCTS Model.

Co-authored-by: lewwang <lwwang@microsoft.com>
This commit is contained in:
Lewen Wang
2021-07-04 16:45:05 +08:00
committed by GitHub
parent 2d4f0e80f9
commit ace7484304
10 changed files with 117 additions and 156 deletions
Download Patch File Download Diff File Expand all files Collapse all files

36
qlib/contrib/model/pytorch_gats.py
View File

@@ -53,7 +53,6 @@ class GATs(Model):
early_stop=20,
loss="mse",
base_model="GRU",
with_pretrain=True,
model_path=None,
optimizer="adam",
GPU=0,
@@ -76,7 +75,6 @@ class GATs(Model):
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() and GPU >= 0 else "cpu")
self.seed = seed
@@ -94,7 +92,6 @@ class GATs(Model):
"\noptimizer : {}"
"\nloss_type : {}"
"\nbase_model : {}"
"\nwith_pretrain : {}"
"\nmodel_path : {}"
"\ndevice : {}"
"\nuse_GPU : {}"
@@ -110,7 +107,6 @@ class GATs(Model):
optimizer.lower(),
loss,
base_model,
with_pretrain,
model_path,
self.device,
self.use_gpu,
@@ -253,24 +249,22 @@ class GATs(Model):
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)
if self.base_model == "LSTM":
pretrained_model = LSTMModel()
elif self.base_model == "GRU":
pretrained_model = GRUModel()
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...")
if self.model_path is not None:
self.logger.info("Loading pretrained model...")
pretrained_model.load_state_dict(torch.load(self.model_path))
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...")