Device dmatrix (#5420)

tests/python/regression_test_utilities.py

@@ -84,7 +84,8 @@ def get_weights_regression(min_weight, max_weight):
     return X, y, w
 
 
-def train_dataset(dataset, param_in, num_rounds=10, scale_features=False):
+def train_dataset(dataset, param_in, num_rounds=10, scale_features=False, DMatrixT=xgb.DMatrix,
+                  dmatrix_params={}):
     param = param_in.copy()
     param["objective"] = dataset.objective
     if dataset.objective == "multi:softmax":
@@ -99,10 +100,13 @@ def train_dataset(dataset, param_in, num_rounds=10, scale_features=False):
     if dataset.use_external_memory:
         np.savetxt('tmptmp_1234.csv', np.hstack((dataset.y.reshape(len(dataset.y), 1), X)),
                    delimiter=',')
-        dtrain = xgb.DMatrix('tmptmp_1234.csv?format=csv&label_column=0#tmptmp_',
+        dtrain = DMatrixT('tmptmp_1234.csv?format=csv&label_column=0#tmptmp_',
                              weight=dataset.w)
+    elif DMatrixT is xgb.DeviceQuantileDMatrix:
+        import cupy as cp
+        dtrain = DMatrixT(cp.array(X), dataset.y, weight=dataset.w, **dmatrix_params)
     else:
-        dtrain = xgb.DMatrix(X, dataset.y, weight=dataset.w)
+        dtrain = DMatrixT(X, dataset.y, weight=dataset.w, **dmatrix_params)
 
     print("Training on dataset: " + dataset.name, file=sys.stderr)
     print("Using parameters: " + str(param), file=sys.stderr)
@@ -139,7 +143,8 @@ def parameter_combinations(variable_param):
     return result
 
 
-def run_suite(param, num_rounds=10, select_datasets=None, scale_features=False):
+def run_suite(param, num_rounds=10, select_datasets=None, scale_features=False,
+              DMatrixT=xgb.DMatrix, dmatrix_params={}):
     """
     Run the given parameters on a range of datasets. Objective and eval metric will be automatically set
     """
@@ -162,7 +167,8 @@ def run_suite(param, num_rounds=10, select_datasets=None, scale_features=False):
     for d in datasets:
         if select_datasets is None or d.name in select_datasets:
             results.append(
-                train_dataset(d, param, num_rounds=num_rounds, scale_features=scale_features))
+                train_dataset(d, param, num_rounds=num_rounds, scale_features=scale_features,
+                              DMatrixT=DMatrixT, dmatrix_params=dmatrix_params))
     return results