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

Fix GPU

Browse Source
This commit is contained in:
Jactus
2020-12-02 13:25:29 +08:00
parent d109d3d44e
commit 7f385345bb
5 changed files with 37 additions and 112 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -83,7 +83,7 @@ class GATs(Model):
self.base_model = base_model
self.with_pretrain = with_pretrain
self.model_path = model_path
self.visible_GPU = GPU
self.device = "cuda:%d" % (GPU) if torch.cuda.is_available() else "cpu"
self.use_gpu = torch.cuda.is_available()
self.seed = seed
@@ -143,11 +143,7 @@ class GATs(Model):
raise NotImplementedError("optimizer {} is not supported!".format(optimizer))
self._fitted = False
if self.use_gpu:
self.GAT_model.cuda()
# set the visible GPU
if self.visible_GPU:
os.environ["CUDA_VISIBLE_DEVICES"] = str(self.visible_GPU)
self.GAT_model.to(self.device)
def mse(self, pred, label):
loss = (pred - label) ** 2
@@ -193,12 +189,8 @@ class GATs(Model):
for idx, count in zip(daily_index, daily_count):
batch = slice(idx, idx + count)
feature = torch.from_numpy(x_train_values[batch]).float()
label = torch.from_numpy(y_train_values[batch]).float()
if self.use_gpu:
feature = feature.cuda()
label = label.cuda()
feature = torch.from_numpy(x_train_values[batch]).float().to(self.device)
label = torch.from_numpy(y_train_values[batch]).float().to(self.device)
pred = self.GAT_model(feature)
loss = self.loss_fn(pred, label)
@@ -224,12 +216,8 @@ class GATs(Model):
for idx, count in zip(daily_index, daily_count):
batch = slice(idx, idx + count)
feature = torch.from_numpy(x_values[batch]).float()
label = torch.from_numpy(y_values[batch]).float()
if self.use_gpu:
feature = feature.cuda()
label = label.cuda()
feature = torch.from_numpy(x_values[batch]).float().to(self.device)
label = torch.from_numpy(y_values[batch]).float().to(self.device)
pred = self.GAT_model(feature)
loss = self.loss_fn(pred, label)
@@ -333,10 +321,7 @@ class GATs(Model):
for idx, count in zip(daily_index, daily_count):
batch = slice(idx, idx + count)
x_batch = torch.from_numpy(x_values[batch]).float()
if self.use_gpu:
x_batch = x_batch.cuda()
x_batch = torch.from_numpy(x_values[batch]).float().to(self.device)
with torch.no_grad():
if self.use_gpu: