From 7f385345bb41d4af7dc6c27265ddc52d3cc8b633 Mon Sep 17 00:00:00 2001 From: Jactus Date: Wed, 2 Dec 2020 13:25:29 +0800 Subject: [PATCH] Fix GPU --- qlib/contrib/model/pytorch_alstm.py | 29 ++++++------------------- qlib/contrib/model/pytorch_gats.py | 29 ++++++------------------- qlib/contrib/model/pytorch_gru.py | 29 ++++++------------------- qlib/contrib/model/pytorch_lstm.py | 29 ++++++------------------- qlib/contrib/model/pytorch_nn.py | 33 ++++++++--------------------- 5 files changed, 37 insertions(+), 112 deletions(-) diff --git a/qlib/contrib/model/pytorch_alstm.py b/qlib/contrib/model/pytorch_alstm.py index 9b5db8042..cfe629fa8 100644 --- a/qlib/contrib/model/pytorch_alstm.py +++ b/qlib/contrib/model/pytorch_alstm.py @@ -76,7 +76,7 @@ class ALSTM(Model): self.early_stop = early_stop self.optimizer = optimizer.lower() self.loss = loss - 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 @@ -131,11 +131,7 @@ class ALSTM(Model): raise NotImplementedError("optimizer {} is not supported!".format(optimizer)) self._fitted = False - if self.use_gpu: - self.ALSTM_model.cuda() - # set the visible GPU - if self.visible_GPU: - os.environ["CUDA_VISIBLE_DEVICES"] = str(self.visible_GPU) + self.ALSTM_model.to(self.device) def mse(self, pred, label): loss = (pred - label) ** 2 @@ -173,12 +169,8 @@ class ALSTM(Model): if len(indices) - i < self.batch_size: break - feature = torch.from_numpy(x_train_values[indices[i : i + self.batch_size]]).float() - label = torch.from_numpy(y_train_values[indices[i : i + self.batch_size]]).float() - - if self.use_gpu: - feature = feature.cuda() - label = label.cuda() + feature = torch.from_numpy(x_train_values[indices[i : i + self.batch_size]]).float().to(self.device) + label = torch.from_numpy(y_train_values[indices[i : i + self.batch_size]]).float().to(self.device) pred = self.ALSTM_model(feature) loss = self.loss_fn(pred, label) @@ -206,12 +198,8 @@ class ALSTM(Model): if len(indices) - i < self.batch_size: break - feature = torch.from_numpy(x_values[indices[i : i + self.batch_size]]).float() - label = torch.from_numpy(y_values[indices[i : i + self.batch_size]]).float() - - if self.use_gpu: - feature = feature.cuda() - label = label.cuda() + feature = torch.from_numpy(x_values[indices[i : i + self.batch_size]]).float().to(self.device) + label = torch.from_numpy(y_values[indices[i : i + self.batch_size]]).float().to(self.device) pred = self.ALSTM_model(feature) loss = self.loss_fn(pred, label) @@ -299,10 +287,7 @@ class ALSTM(Model): else: end = begin + self.batch_size - x_batch = torch.from_numpy(x_values[begin:end]).float() - - if self.use_gpu: - x_batch = x_batch.cuda() + x_batch = torch.from_numpy(x_values[begin:end]).float().to(self.device) with torch.no_grad(): if self.use_gpu: diff --git a/qlib/contrib/model/pytorch_gats.py b/qlib/contrib/model/pytorch_gats.py index 5bc91eb3b..1759ae5ba 100644 --- a/qlib/contrib/model/pytorch_gats.py +++ b/qlib/contrib/model/pytorch_gats.py @@ -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: diff --git a/qlib/contrib/model/pytorch_gru.py b/qlib/contrib/model/pytorch_gru.py index 7974e7c33..8d5a9e461 100755 --- a/qlib/contrib/model/pytorch_gru.py +++ b/qlib/contrib/model/pytorch_gru.py @@ -76,7 +76,7 @@ class GRU(Model): self.early_stop = early_stop self.optimizer = optimizer.lower() self.loss = loss - 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 @@ -131,11 +131,7 @@ class GRU(Model): raise NotImplementedError("optimizer {} is not supported!".format(optimizer)) self._fitted = False - if self.use_gpu: - self.gru_model.cuda() - # set the visible GPU - if self.visible_GPU: - os.environ["CUDA_VISIBLE_DEVICES"] = str(self.visible_GPU) + self.gru_model.to(self.device) def mse(self, pred, label): loss = (pred - label) ** 2 @@ -173,12 +169,8 @@ class GRU(Model): if len(indices) - i < self.batch_size: break - feature = torch.from_numpy(x_train_values[indices[i : i + self.batch_size]]).float() - label = torch.from_numpy(y_train_values[indices[i : i + self.batch_size]]).float() - - if self.use_gpu: - feature = feature.cuda() - label = label.cuda() + feature = torch.from_numpy(x_train_values[indices[i : i + self.batch_size]]).float().to(self.device) + label = torch.from_numpy(y_train_values[indices[i : i + self.batch_size]]).float().to(self.device) pred = self.gru_model(feature) loss = self.loss_fn(pred, label) @@ -206,12 +198,8 @@ class GRU(Model): if len(indices) - i < self.batch_size: break - feature = torch.from_numpy(x_values[indices[i : i + self.batch_size]]).float() - label = torch.from_numpy(y_values[indices[i : i + self.batch_size]]).float() - - if self.use_gpu: - feature = feature.cuda() - label = label.cuda() + feature = torch.from_numpy(x_values[indices[i : i + self.batch_size]]).float().to(self.device) + label = torch.from_numpy(y_values[indices[i : i + self.batch_size]]).float().to(self.device) pred = self.gru_model(feature) loss = self.loss_fn(pred, label) @@ -299,10 +287,7 @@ class GRU(Model): else: end = begin + self.batch_size - x_batch = torch.from_numpy(x_values[begin:end]).float() - - if self.use_gpu: - x_batch = x_batch.cuda() + x_batch = torch.from_numpy(x_values[begin:end]).float().to(self.device) with torch.no_grad(): if self.use_gpu: diff --git a/qlib/contrib/model/pytorch_lstm.py b/qlib/contrib/model/pytorch_lstm.py index 2eb6ae69f..8032c54eb 100755 --- a/qlib/contrib/model/pytorch_lstm.py +++ b/qlib/contrib/model/pytorch_lstm.py @@ -76,7 +76,7 @@ class LSTM(Model): self.early_stop = early_stop self.optimizer = optimizer.lower() self.loss = loss - 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 @@ -131,11 +131,7 @@ class LSTM(Model): raise NotImplementedError("optimizer {} is not supported!".format(optimizer)) self._fitted = False - if self.use_gpu: - self.lstm_model.cuda() - # set the visible GPU - if self.visible_GPU: - os.environ["CUDA_VISIBLE_DEVICES"] = str(self.visible_GPU) + self.lstm_model.to(self.device) def mse(self, pred, label): loss = (pred - label) ** 2 @@ -173,12 +169,8 @@ class LSTM(Model): if len(indices) - i < self.batch_size: break - feature = torch.from_numpy(x_train_values[indices[i : i + self.batch_size]]).float() - label = torch.from_numpy(y_train_values[indices[i : i + self.batch_size]]).float() - - if self.use_gpu: - feature = feature.cuda() - label = label.cuda() + feature = torch.from_numpy(x_train_values[indices[i : i + self.batch_size]]).float().to(self.device) + label = torch.from_numpy(y_train_values[indices[i : i + self.batch_size]]).float().to(self.device) pred = self.lstm_model(feature) loss = self.loss_fn(pred, label) @@ -206,12 +198,8 @@ class LSTM(Model): if len(indices) - i < self.batch_size: break - feature = torch.from_numpy(x_values[indices[i : i + self.batch_size]]).float() - label = torch.from_numpy(y_values[indices[i : i + self.batch_size]]).float() - - if self.use_gpu: - feature = feature.cuda() - label = label.cuda() + feature = torch.from_numpy(x_values[indices[i : i + self.batch_size]]).float().to(self.device) + label = torch.from_numpy(y_values[indices[i : i + self.batch_size]]).float().to(self.device) pred = self.lstm_model(feature) loss = self.loss_fn(pred, label) @@ -299,10 +287,7 @@ class LSTM(Model): else: end = begin + self.batch_size - x_batch = torch.from_numpy(x_values[begin:end]).float() - - if self.use_gpu: - x_batch = x_batch.cuda() + x_batch = torch.from_numpy(x_values[begin:end]).float().to(self.device) with torch.no_grad(): if self.use_gpu: diff --git a/qlib/contrib/model/pytorch_nn.py b/qlib/contrib/model/pytorch_nn.py index a9d8a38d0..ed4ba8692 100644 --- a/qlib/contrib/model/pytorch_nn.py +++ b/qlib/contrib/model/pytorch_nn.py @@ -79,7 +79,7 @@ class DNNModelPytorch(Model): self.lr_decay_steps = lr_decay_steps self.optimizer = optimizer.lower() self.loss_type = loss - 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 @@ -147,11 +147,7 @@ class DNNModelPytorch(Model): ) self._fitted = False - if self.use_GPU: - self.dnn_model.cuda() - # set the visible GPU - if self.visible_GPU: - os.environ["CUDA_VISIBLE_DEVICES"] = str(self.visible_GPU) + self.dnn_model.to(self.device) def fit( self, @@ -188,13 +184,9 @@ class DNNModelPytorch(Model): w_train_values = torch.from_numpy(w_train.values).float() train_num = y_train_values.shape[0] # prepare validation data - x_val_auto = torch.from_numpy(x_valid.values).float() - y_val_auto = torch.from_numpy(y_valid.values).float() - w_val_auto = torch.from_numpy(w_valid.values).float() - if self.use_GPU: - x_val_auto = x_val_auto.cuda() - y_val_auto = y_val_auto.cuda() - w_val_auto = w_val_auto.cuda() + x_val_auto = torch.from_numpy(x_valid.values).float().to(self.device) + y_val_auto = torch.from_numpy(y_valid.values).float().to(self.device) + w_val_auto = torch.from_numpy(w_valid.values).float().to(self.device) for step in range(self.max_steps): if stop_steps >= self.early_stop_rounds: @@ -205,14 +197,9 @@ class DNNModelPytorch(Model): self.dnn_model.train() self.train_optimizer.zero_grad() choice = np.random.choice(train_num, self.batch_size) - x_batch_auto = x_train_values[choice] - y_batch_auto = y_train_values[choice] - w_batch_auto = w_train_values[choice] - - if self.use_GPU: - x_batch_auto = x_batch_auto.cuda() - y_batch_auto = y_batch_auto.cuda() - w_batch_auto = w_batch_auto.cuda() + x_batch_auto = x_train_values[choice].to(self.device) + y_batch_auto = y_train_values[choice].to(self.device) + w_batch_auto = w_train_values[choice].to(self.device) # forward preds = self.dnn_model(x_batch_auto) @@ -277,9 +264,7 @@ class DNNModelPytorch(Model): if not self._fitted: raise ValueError("model is not fitted yet!") x_test_pd = dataset.prepare("test", col_set="feature") - x_test = torch.from_numpy(x_test_pd.values).float() - if self.use_GPU: - x_test = x_test.cuda() + x_test = torch.from_numpy(x_test_pd.values).float().to(self.device) self.dnn_model.eval() with torch.no_grad():