[SPARK-36425][PYSPARK][ML] Support CrossValidatorModel get standard deviation of metrics for each paramMap

Signed-off-by: Weichen Xu <weichen.xudatabricks.com>

### What changes were proposed in this pull request?
Support CrossValidatorModel get standard deviation of metrics for each paramMap.

### Why are the changes needed?
So that in mlflow autologging, we can log standard deviation of metrics which is useful.

### Does this PR introduce _any_ user-facing change?
Yes.
`CrossValidatorModel` add a public attribute `stdMetrics` which are the standard deviation of metrics for each paramMap

### How was this patch tested?
Unit test.

Closes #33652 from WeichenXu123/add_std_metric.

Authored-by: Weichen Xu <[email protected]>
Signed-off-by: Hyukjin Kwon <[email protected]>

Author: WeichenXu123

python/pyspark/ml/tests/test_tuning.py

@@ -194,6 +194,18 @@ def assert_param_maps_equal(self, paramMaps1, paramMaps2):
 
 class CrossValidatorTests(SparkSessionTestCase, ValidatorTestUtilsMixin):
 
+    def test_gen_avg_and_std_metrics(self):
+        metrics_all = [
+            [1.0, 3.0, 2.0, 4.0],
+            [3.0, 2.0, 2.0, 4.0],
+            [3.0, 2.5, 2.1, 8.0],
+        ]
+        avg_metrics, std_metrics = CrossValidator._gen_avg_and_std_metrics(metrics_all)
+        assert np.allclose(avg_metrics, [2.33333333, 2.5, 2.03333333, 5.33333333])
+        assert np.allclose(std_metrics, [0.94280904, 0.40824829, 0.04714045, 1.88561808])
+        assert isinstance(avg_metrics, list)
+        assert isinstance(std_metrics, list)
+
     def test_copy(self):
         dataset = self.spark.createDataFrame([
             (10, 10.0),
@@ -232,6 +244,7 @@ def test_copy(self):
         for index in range(len(cvModel.avgMetrics)):
             self.assertTrue(abs(cvModel.avgMetrics[index] - cvModelCopied.avgMetrics[index])
                             < 0.0001)
+        self.assertTrue(np.allclose(cvModel.stdMetrics, cvModelCopied.stdMetrics))
         # SPARK-32092: CrossValidatorModel.copy() needs to copy all existing params
         for param in [
             lambda x: x.getNumFolds(),
@@ -246,6 +259,12 @@ def test_copy(self):
             'foo',
             "Changing the original avgMetrics should not affect the copied model"
         )
+        cvModel.stdMetrics[0] = 'foo'
+        self.assertNotEqual(
+            cvModelCopied.stdMetrics[0],
+            'foo',
+            "Changing the original stdMetrics should not affect the copied model"
+        )
         cvModel.subModels[0][0].getInducedError = lambda: 'foo'
         self.assertNotEqual(
             cvModelCopied.subModels[0][0].getInducedError(),
@@ -353,6 +372,15 @@ def _run_test_save_load_trained_model(self, LogisticRegressionCls, LogisticRegre
             loadedCvModel.isSet(param) for param in loadedCvModel.params
         ))
 
+        # mimic old version CrossValidatorModel (without stdMetrics attribute)
+        # test loading model backwards compatibility
+        cvModel2 = cvModel.copy()
+        cvModel2.stdMetrics = []
+        cvModelPath2 = temp_path + "/cvModel2"
+        cvModel2.save(cvModelPath2)
+        loadedCvModel2 = CrossValidatorModel.load(cvModelPath2)
+        assert loadedCvModel2.stdMetrics == []
+
     def test_save_load_trained_model(self):
         self._run_test_save_load_trained_model(LogisticRegression, LogisticRegressionModel)
         self._run_test_save_load_trained_model(DummyLogisticRegression,
@@ -414,6 +442,7 @@ def test_parallel_evaluation(self):
         cv.setParallelism(2)
         cvParallelModel = cv.fit(dataset)
         self.assertEqual(cvSerialModel.avgMetrics, cvParallelModel.avgMetrics)
+        self.assertEqual(cvSerialModel.stdMetrics, cvParallelModel.stdMetrics)
 
     def test_expose_sub_models(self):
         temp_path = tempfile.mkdtemp()

python/pyspark/ml/tuning.py

@@ -499,6 +499,10 @@ def load(self, path):
             bestModelPath = os.path.join(path, 'bestModel')
             bestModel = DefaultParamsReader.loadParamsInstance(bestModelPath, self.sc)
             avgMetrics = metadata['avgMetrics']
+            if 'stdMetrics' in metadata:
+                stdMetrics = metadata['stdMetrics']
+            else:
+                stdMetrics = None
             persistSubModels = ('persistSubModels' in metadata) and metadata['persistSubModels']
 
             if persistSubModels:
@@ -512,7 +516,9 @@ def load(self, path):
             else:
                 subModels = None
 
-            cvModel = CrossValidatorModel(bestModel, avgMetrics=avgMetrics, subModels=subModels)
+            cvModel = CrossValidatorModel(
+                bestModel, avgMetrics=avgMetrics, subModels=subModels, stdMetrics=stdMetrics
+            )
             cvModel = cvModel._resetUid(metadata['uid'])
             cvModel.set(cvModel.estimator, estimator)
             cvModel.set(cvModel.estimatorParamMaps, estimatorParamMaps)
@@ -536,6 +542,9 @@ def saveImpl(self, path):
             .getValidatorModelWriterPersistSubModelsParam(self)
         extraMetadata = {'avgMetrics': instance.avgMetrics,
                          'persistSubModels': persistSubModels}
+        if instance.stdMetrics:
+            extraMetadata['stdMetrics'] = instance.stdMetrics
+
         _ValidatorSharedReadWrite.saveImpl(path, instance, self.sc, extraMetadata=extraMetadata)
         bestModelPath = os.path.join(path, 'bestModel')
         instance.bestModel.save(bestModelPath)
@@ -710,13 +719,19 @@ def setCollectSubModels(self, value):
         """
         return self._set(collectSubModels=value)
 
+    @staticmethod
+    def _gen_avg_and_std_metrics(metrics_all):
+        avg_metrics = np.mean(metrics_all, axis=0)
+        std_metrics = np.std(metrics_all, axis=0)
+        return list(avg_metrics), list(std_metrics)
+
     def _fit(self, dataset):
         est = self.getOrDefault(self.estimator)
         epm = self.getOrDefault(self.estimatorParamMaps)
         numModels = len(epm)
         eva = self.getOrDefault(self.evaluator)
         nFolds = self.getOrDefault(self.numFolds)
-        metrics = [0.0] * numModels
+        metrics_all = [[0.0] * numModels for i in range(nFolds)]
 
         pool = ThreadPool(processes=min(self.getParallelism(), numModels))
         subModels = None
@@ -733,19 +748,21 @@ def _fit(self, dataset):
                 inheritable_thread_target,
                 _parallelFitTasks(est, train, eva, validation, epm, collectSubModelsParam))
             for j, metric, subModel in pool.imap_unordered(lambda f: f(), tasks):
-                metrics[j] += (metric / nFolds)
+                metrics_all[i][j] = metric
                 if collectSubModelsParam:
                     subModels[i][j] = subModel
 
             validation.unpersist()
             train.unpersist()
 
+        metrics, std_metrics = CrossValidator._gen_avg_and_std_metrics(metrics_all)
+
         if eva.isLargerBetter():
             bestIndex = np.argmax(metrics)
         else:
             bestIndex = np.argmin(metrics)
         bestModel = est.fit(dataset, epm[bestIndex])
-        return self._copyValues(CrossValidatorModel(bestModel, metrics, subModels))
+        return self._copyValues(CrossValidatorModel(bestModel, metrics, subModels, std_metrics))
 
     def _kFold(self, dataset):
         nFolds = self.getOrDefault(self.numFolds)
@@ -875,15 +892,20 @@ def _to_java(self):
 
 class CrossValidatorModel(Model, _CrossValidatorParams, MLReadable, MLWritable):
     """
-
     CrossValidatorModel contains the model with the highest average cross-validation
     metric across folds and uses this model to transform input data. CrossValidatorModel
     also tracks the metrics for each param map evaluated.
 
     .. versionadded:: 1.4.0
+
+    Notes
+    -----
+    Since version 3.3.0, CrossValidatorModel contains a new attribute "stdMetrics",
+    which represent standard deviation of metrics for each paramMap in
+    CrossValidator.estimatorParamMaps.
     """
 
-    def __init__(self, bestModel, avgMetrics=None, subModels=None):
+    def __init__(self, bestModel, avgMetrics=None, subModels=None, stdMetrics=None):
         super(CrossValidatorModel, self).__init__()
         #: best model from cross validation
         self.bestModel = bestModel
@@ -892,6 +914,9 @@ def __init__(self, bestModel, avgMetrics=None, subModels=None):
         self.avgMetrics = avgMetrics or []
         #: sub model list from cross validation
         self.subModels = subModels
+        #: standard deviation of metrics for each paramMap in
+        #: CrossValidator.estimatorParamMaps, in the corresponding order.
+        self.stdMetrics = stdMetrics or []
 
     def _transform(self, dataset):
         return self.bestModel.transform(dataset)
@@ -924,7 +949,9 @@ def copy(self, extra=None):
             [sub_model.copy() for sub_model in fold_sub_models]
             for fold_sub_models in self.subModels
         ]
-        return self._copyValues(CrossValidatorModel(bestModel, avgMetrics, subModels), extra=extra)
+        stdMetrics = list(self.stdMetrics)
+        return self._copyValues(CrossValidatorModel(bestModel, avgMetrics, subModels, stdMetrics),
+                                extra=extra)
 
     @since("2.3.0")
     def write(self):