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

Collect all contrib models in __init__ and add unit tests for init

Browse Source
This commit is contained in:
D-X-Y
2021-03-28 10:39:28 +00:00
parent 8a2e7b62af
commit 0386df7b16
17 changed files with 115 additions and 46 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -8,6 +8,7 @@ from __future__ import print_function
import os
import numpy as np
import pandas as pd
from typing import Text, Union
import copy
from ...utils import (
unpack_archive_with_buffer,
@@ -83,7 +84,6 @@ class GATs(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.use_gpu = torch.cuda.is_available()
self.seed = seed
self.logger.info(
@@ -310,11 +310,11 @@ class GATs(Model):
if self.use_gpu:
torch.cuda.empty_cache()
def predict(self, dataset):
def predict(self, dataset: DatasetH, segment: Union[Text, slice] = "test"):
if not self.fitted:
raise ValueError("model is not fitted yet!")
x_test = dataset.prepare("test", col_set="feature")
x_test = dataset.prepare(segment, col_set="feature")
index = x_test.index
self.GAT_model.eval()
x_values = x_test.values