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

Fix code with block.

Browse Source
This commit is contained in:
lwwang1995
2020-11-27 22:44:28 +08:00
parent c5a3b74a96
commit 1353e81b5b
7 changed files with 42 additions and 136 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -126,9 +126,7 @@ class GATs(Model):
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)
)
raise NotImplementedError("optimizer {} is not supported!".format(optimizer))
self._fitted = False
if self.use_gpu:
@@ -258,22 +256,14 @@ class GATs(Model):
self.logger.info("Loading pretrained model...")
if self.base_model == "LSTM":
pretrained_model = LSTMModel()
pretrained_model.load_state_dict(
torch.load("benchmarks/LSTM/model_lstm_csi300.pkl")
)
pretrained_model.load_state_dict(torch.load("benchmarks/LSTM/model_lstm_csi300.pkl"))
elif self.base_model == "GRU":
pretrained_model = GRUModel()
pretrained_model.load_state_dict(
torch.load("benchmarks/GRU/model_gru_csi300.pkl")
)
pretrained_model.load_state_dict(torch.load("benchmarks/GRU/model_gru_csi300.pkl"))
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
}
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...")
@@ -343,9 +333,7 @@ class GATs(Model):
class GATModel(nn.Module):
def __init__(
self, d_feat=6, hidden_size=64, num_layers=2, dropout=0.0, base_model="GRU"
):
def __init__(self, d_feat=6, hidden_size=64, num_layers=2, dropout=0.0, base_model="GRU"):
super().__init__()
if base_model == "GRU":