feat: pytorch benchmarks support mps device

qlib/contrib/model/pytorch_adarnn.py

@@ -12,7 +12,7 @@
 import torch.nn.functional as F
 import torch.optim as optim
 from torch.autograd import Function
-from qlib.contrib.model.pytorch_utils import count_parameters
+from qlib.contrib.model.pytorch_utils import count_parameters, get_device
 from qlib.data.dataset import DatasetH
 from qlib.data.dataset.handler import DataHandlerLP
 from qlib.log import get_module_logger
@@ -81,7 +81,7 @@ def __init__(
         self.optimizer = optimizer.lower()
         self.loss = loss
         self.n_splits = n_splits
-        self.device = torch.device("cuda:%d" % GPU if torch.cuda.is_available() and GPU >= 0 else "cpu")
+        self.device = get_device(GPU)
         self.seed = seed
 
         self.logger.info(
@@ -162,7 +162,9 @@ def train_AdaRNN(self, train_loader_list, epoch, dist_old=None, weight_mat=None)
             list_label = []
             for data in data_all:
                 # feature :[36, 24, 6]
-                feature, label_reg = data[0].to(self.device).float(), data[1].to(self.device).float()
+                feature, label_reg = data[0].to(self.device, dtype=torch.float32), data[1].to(
+                    self.device, dtype=torch.float32
+                )
                 list_feat.append(feature)
                 list_label.append(label_reg)
             flag = False
@@ -396,7 +398,7 @@ def __init__(
         self.model_type = model_type
         self.trans_loss = trans_loss
         self.len_seq = len_seq
-        self.device = torch.device("cuda:%d" % GPU if torch.cuda.is_available() and GPU >= 0 else "cpu")
+        self.device = get_device(GPU)
         in_size = self.n_input
 
         features = nn.ModuleList()
@@ -558,7 +560,7 @@ def __init__(self, loss_type="cosine", input_dim=512, GPU=0):
         """
         self.loss_type = loss_type
         self.input_dim = input_dim
-        self.device = torch.device("cuda:%d" % GPU if torch.cuda.is_available() and GPU >= 0 else "cpu")
+        self.device = get_device(GPU)
 
     def compute(self, X, Y):
         """Compute adaptation loss

qlib/contrib/model/pytorch_add.py

@@ -17,7 +17,7 @@
 import torch.optim as optim
 from qlib.contrib.model.pytorch_gru import GRUModel
 from qlib.contrib.model.pytorch_lstm import LSTMModel
-from qlib.contrib.model.pytorch_utils import count_parameters
+from qlib.contrib.model.pytorch_utils import count_parameters, get_device
 from qlib.data.dataset import DatasetH
 from qlib.data.dataset.handler import DataHandlerLP
 from qlib.log import get_module_logger
@@ -83,7 +83,7 @@ def __init__(
         self.optimizer = optimizer.lower()
         self.base_model = base_model
         self.model_path = model_path
-        self.device = torch.device("cuda:%d" % (GPU) if torch.cuda.is_available() and GPU >= 0 else "cpu")
+        self.device = get_device(GPU)
         self.seed = seed
 
         self.gamma = gamma

qlib/contrib/model/pytorch_alstm.py

@@ -16,7 +16,7 @@
 import torch.nn as nn
 import torch.optim as optim
 
-from .pytorch_utils import count_parameters
+from .pytorch_utils import count_parameters, get_device
 from ...model.base import Model
 from ...data.dataset import DatasetH
 from ...data.dataset.handler import DataHandlerLP
@@ -70,7 +70,7 @@ def __init__(
         self.early_stop = early_stop
         self.optimizer = optimizer.lower()
         self.loss = loss
-        self.device = torch.device("cuda:%d" % (GPU) if torch.cuda.is_available() and GPU >= 0 else "cpu")
+        self.device = get_device(GPU)
         self.seed = seed
 
         self.logger.info(

qlib/contrib/model/pytorch_alstm_ts.py

@@ -17,7 +17,7 @@
 import torch.optim as optim
 from torch.utils.data import DataLoader
 
-from .pytorch_utils import count_parameters
+from .pytorch_utils import count_parameters, get_device
 from ...model.base import Model
 from ...data.dataset import DatasetH
 from ...data.dataset.handler import DataHandlerLP
@@ -74,7 +74,7 @@ def __init__(
         self.early_stop = early_stop
         self.optimizer = optimizer.lower()
         self.loss = loss
-        self.device = torch.device("cuda:%d" % (GPU) if torch.cuda.is_available() and GPU >= 0 else "cpu")
+        self.device = get_device(GPU)
         self.n_jobs = n_jobs
         self.seed = seed
 
@@ -171,11 +171,11 @@ def train_epoch(self, data_loader):
         self.ALSTM_model.train()
 
         for data, weight in data_loader:
-            feature = data[:, :, 0:-1].to(self.device)
-            label = data[:, -1, -1].to(self.device)
+            feature = data[:, :, 0:-1].to(self.device, dtype=torch.float32)
+            label = data[:, -1, -1].to(self.device, dtype=torch.float32)
 
-            pred = self.ALSTM_model(feature.float())
-            loss = self.loss_fn(pred, label, weight.to(self.device))
+            pred = self.ALSTM_model(feature)
+            loss = self.loss_fn(pred, label, weight.to(self.device, dtype=torch.float32))
 
             self.train_optimizer.zero_grad()
             loss.backward()
@@ -189,13 +189,13 @@ def test_epoch(self, data_loader):
         losses = []
 
         for data, weight in data_loader:
-            feature = data[:, :, 0:-1].to(self.device)
+            feature = data[:, :, 0:-1].to(self.device, dtype=torch.float32)
             # feature[torch.isnan(feature)] = 0
-            label = data[:, -1, -1].to(self.device)
+            label = data[:, -1, -1].to(self.device, dtype=torch.float32)
 
             with torch.no_grad():
-                pred = self.ALSTM_model(feature.float())
-                loss = self.loss_fn(pred, label, weight.to(self.device))
+                pred = self.ALSTM_model(feature)
+                loss = self.loss_fn(pred, label, weight.to(self.device, dtype=torch.float32))
                 losses.append(loss.item())
 
                 score = self.metric_fn(pred, label)
@@ -295,10 +295,10 @@ def predict(self, dataset: DatasetH, segment: Union[Text, slice] = "test"):
         preds = []
 
         for data in test_loader:
-            feature = data[:, :, 0:-1].to(self.device)
+            feature = data[:, :, 0:-1].to(self.device, dtype=torch.float32)
 
             with torch.no_grad():
-                pred = self.ALSTM_model(feature.float()).detach().cpu().numpy()
+                pred = self.ALSTM_model(feature).detach().cpu().numpy()
 
             preds.append(pred)
 

qlib/contrib/model/pytorch_gats.py

@@ -15,7 +15,7 @@
 import torch.nn as nn
 import torch.optim as optim
 
-from .pytorch_utils import count_parameters
+from .pytorch_utils import count_parameters, get_device
 from ...model.base import Model
 from ...data.dataset import DatasetH
 from ...data.dataset.handler import DataHandlerLP
@@ -75,7 +75,7 @@ def __init__(
         self.loss = loss
         self.base_model = base_model
         self.model_path = model_path
-        self.device = torch.device("cuda:%d" % (GPU) if torch.cuda.is_available() and GPU >= 0 else "cpu")
+        self.device = get_device(GPU)
         self.seed = seed
 
         self.logger.info(

qlib/contrib/model/pytorch_gats_ts.py

@@ -16,7 +16,7 @@
 from torch.utils.data import DataLoader
 from torch.utils.data import Sampler
 
-from .pytorch_utils import count_parameters
+from .pytorch_utils import count_parameters, get_device
 from ...model.base import Model
 from ...data.dataset.handler import DataHandlerLP
 from ...contrib.model.pytorch_lstm import LSTMModel
@@ -94,7 +94,7 @@ def __init__(
         self.loss = loss
         self.base_model = base_model
         self.model_path = model_path
-        self.device = torch.device("cuda:%d" % (GPU) if torch.cuda.is_available() and GPU >= 0 else "cpu")
+        self.device = get_device(GPU)
         self.n_jobs = n_jobs
         self.seed = seed
 
@@ -198,10 +198,10 @@ def train_epoch(self, data_loader):
 
         for data in data_loader:
             data = data.squeeze()
-            feature = data[:, :, 0:-1].to(self.device)
-            label = data[:, -1, -1].to(self.device)
+            feature = data[:, :, 0:-1].to(self.device, dtype=torch.float32)
+            label = data[:, -1, -1].to(self.device, dtype=torch.float32)
 
-            pred = self.GAT_model(feature.float())
+            pred = self.GAT_model(feature)
             loss = self.loss_fn(pred, label)
 
             self.train_optimizer.zero_grad()
@@ -217,11 +217,11 @@ def test_epoch(self, data_loader):
 
         for data in data_loader:
             data = data.squeeze()
-            feature = data[:, :, 0:-1].to(self.device)
+            feature = data[:, :, 0:-1].to(self.device, dtype=torch.float32)
             # feature[torch.isnan(feature)] = 0
-            label = data[:, -1, -1].to(self.device)
+            label = data[:, -1, -1].to(self.device, dtype=torch.float32)
 
-            pred = self.GAT_model(feature.float())
+            pred = self.GAT_model(feature)
             loss = self.loss_fn(pred, label)
             losses.append(loss.item())
 
@@ -325,10 +325,10 @@ def predict(self, dataset):
 
         for data in test_loader:
             data = data.squeeze()
-            feature = data[:, :, 0:-1].to(self.device)
+            feature = data[:, :, 0:-1].to(self.device, dtype=torch.float32)
 
             with torch.no_grad():
-                pred = self.GAT_model(feature.float()).detach().cpu().numpy()
+                pred = self.GAT_model(feature).detach().cpu().numpy()
 
             preds.append(pred)
 

qlib/contrib/model/pytorch_general_nn.py

@@ -17,7 +17,7 @@
 
 from qlib.data.dataset.weight import Reweighter
 
-from .pytorch_utils import count_parameters
+from .pytorch_utils import count_parameters, get_device
 from ...model.base import Model
 from ...data.dataset import DatasetH, TSDatasetH
 from ...data.dataset.handler import DataHandlerLP
@@ -83,7 +83,7 @@ def __init__(
         self.optimizer = optimizer.lower()
         self.loss = loss
         self.weight_decay = weight_decay
-        self.device = torch.device("cuda:%d" % (GPU) if torch.cuda.is_available() and GPU >= 0 else "cpu")
+        self.device = get_device(GPU)
         self.n_jobs = n_jobs
         self.seed = seed
 
@@ -189,12 +189,12 @@ def _get_fl(self, data: torch.Tensor):
         """
         if data.dim() == 3:
             # it is a time series dataset
-            feature = data[:, :, 0:-1].to(self.device)
-            label = data[:, -1, -1].to(self.device)
+            feature = data[:, :, 0:-1].to(self.device, dtype=torch.float32)
+            label = data[:, -1, -1].to(self.device, dtype=torch.float32)
         elif data.dim() == 2:
             # it is a tabular dataset
-            feature = data[:, 0:-1].to(self.device)
-            label = data[:, -1].to(self.device)
+            feature = data[:, 0:-1].to(self.device, dtype=torch.float32)
+            label = data[:, -1].to(self.device, dtype=torch.float32)
         else:
             raise ValueError("Unsupported data shape.")
         return feature, label
@@ -205,8 +205,8 @@ def train_epoch(self, data_loader):
         for data, weight in data_loader:
             feature, label = self._get_fl(data)
 
-            pred = self.dnn_model(feature.float())
-            loss = self.loss_fn(pred, label, weight.to(self.device))
+            pred = self.dnn_model(feature)
+            loss = self.loss_fn(pred, label, weight.to(self.device, dtype=torch.float32))
 
             self.train_optimizer.zero_grad()
             loss.backward()
@@ -223,8 +223,8 @@ def test_epoch(self, data_loader):
             feature, label = self._get_fl(data)
 
             with torch.no_grad():
-                pred = self.dnn_model(feature.float())
-                loss = self.loss_fn(pred, label, weight.to(self.device))
+                pred = self.dnn_model(feature)
+                loss = self.loss_fn(pred, label, weight.to(self.device, dtype=torch.float32))
                 losses.append(loss.item())
 
                 score = self.metric_fn(pred, label)
@@ -357,10 +357,10 @@ def predict(
 
         for data in test_loader:
             feature, _ = self._get_fl(data)
-            feature = feature.to(self.device)
+            feature = feature.to(self.device, dtype=torch.float32)
 
             with torch.no_grad():
-                pred = self.dnn_model(feature.float()).detach().cpu().numpy()
+                pred = self.dnn_model(feature).detach().cpu().numpy()
 
             preds.append(pred)
 

qlib/contrib/model/pytorch_gru.py

@@ -19,7 +19,7 @@
 from ...log import get_module_logger
 from ...model.base import Model
 from ...utils import get_or_create_path
-from .pytorch_utils import count_parameters
+from .pytorch_utils import count_parameters, get_device
 
 
 class GRU(Model):
@@ -70,7 +70,7 @@ def __init__(
         self.early_stop = early_stop
         self.optimizer = optimizer.lower()
         self.loss = loss
-        self.device = torch.device("cuda:%d" % (GPU) if torch.cuda.is_available() and GPU >= 0 else "cpu")
+        self.device = get_device(GPU)
         self.seed = seed
 
         self.logger.info(

qlib/contrib/model/pytorch_gru_ts.py

@@ -16,7 +16,7 @@
 import torch.optim as optim
 from torch.utils.data import DataLoader
 
-from .pytorch_utils import count_parameters
+from .pytorch_utils import count_parameters, get_device
 from ...model.base import Model
 from ...data.dataset.handler import DataHandlerLP
 from ...model.utils import ConcatDataset
@@ -72,7 +72,7 @@ def __init__(
         self.early_stop = early_stop
         self.optimizer = optimizer.lower()
         self.loss = loss
-        self.device = torch.device("cuda:%d" % (GPU) if torch.cuda.is_available() and GPU >= 0 else "cpu")
+        self.device = get_device(GPU)
         self.n_jobs = n_jobs
         self.seed = seed
 
@@ -165,11 +165,11 @@ def train_epoch(self, data_loader):
         self.GRU_model.train()
 
         for data, weight in data_loader:
-            feature = data[:, :, 0:-1].to(self.device)
-            label = data[:, -1, -1].to(self.device)
+            feature = data[:, :, 0:-1].to(self.device, dtype=torch.float32)
+            label = data[:, -1, -1].to(self.device, dtype=torch.float32)
 
-            pred = self.GRU_model(feature.float())
-            loss = self.loss_fn(pred, label, weight.to(self.device))
+            pred = self.GRU_model(feature)
+            loss = self.loss_fn(pred, label, weight.to(self.device, dtype=torch.float32))
 
             self.train_optimizer.zero_grad()
             loss.backward()
@@ -183,13 +183,13 @@ def test_epoch(self, data_loader):
         losses = []
 
         for data, weight in data_loader:
-            feature = data[:, :, 0:-1].to(self.device)
+            feature = data[:, :, 0:-1].to(self.device, dtype=torch.float32)
             # feature[torch.isnan(feature)] = 0
-            label = data[:, -1, -1].to(self.device)
+            label = data[:, -1, -1].to(self.device, dtype=torch.float32)
 
             with torch.no_grad():
-                pred = self.GRU_model(feature.float())
-                loss = self.loss_fn(pred, label, weight.to(self.device))
+                pred = self.GRU_model(feature)
+                loss = self.loss_fn(pred, label, weight.to(self.device, dtype=torch.float32))
                 losses.append(loss.item())
 
                 score = self.metric_fn(pred, label)
@@ -289,10 +289,10 @@ def predict(self, dataset):
         preds = []
 
         for data in test_loader:
-            feature = data[:, :, 0:-1].to(self.device)
+            feature = data[:, :, 0:-1].to(self.device, dtype=torch.float32)
 
             with torch.no_grad():
-                pred = self.GRU_model(feature.float()).detach().cpu().numpy()
+                pred = self.GRU_model(feature).detach().cpu().numpy()
 
             preds.append(pred)
 

qlib/contrib/model/pytorch_hist.py

@@ -16,7 +16,7 @@
 import torch
 import torch.nn as nn
 import torch.optim as optim
-from .pytorch_utils import count_parameters
+from .pytorch_utils import count_parameters, get_device
 from ...model.base import Model
 from ...data.dataset import DatasetH
 from ...data.dataset.handler import DataHandlerLP
@@ -80,7 +80,7 @@ def __init__(
         self.model_path = model_path
         self.stock2concept = stock2concept
         self.stock_index = stock_index
-        self.device = torch.device("cuda:%d" % (GPU) if torch.cuda.is_available() and GPU >= 0 else "cpu")
+        self.device = get_device(GPU)
         self.seed = seed
 
         self.logger.info(