Split Evaluate on test task into two tasks

Co-Authored-By: Adam Schill Collberg<[email protected]>
Co-Authored-By: Mats Rydberg <[email protected]>

Author: breakanalysis

pipeline/src/main/java/org/neo4j/gds/ml/pipeline/nodePipeline/classification/train/NodeClassificationTrain.java

@@ -159,6 +159,7 @@ public static List<Task> progressTasks(NodePropertyPredictionSplitConfig splitCo
                 numberOfModelSelectionTrials
             ),
             ClassifierTrainer.progressTask("Train best model", 5 * trainSetSize),
+            Tasks.leaf("Evaluate on train data", trainSetSize),
             Tasks.leaf("Evaluate on test data", testSetSize),
             ClassifierTrainer.progressTask("Retrain best model", 5 * nodeCount)
         );
@@ -401,16 +402,18 @@ private void evaluateBestModel(
         );
         progressTracker.endSubTask("Train best model");
 
-        progressTracker.beginSubTask("Evaluate on test data");
-        progressTracker.setSteps(outerSplit.testSet().size() + outerSplit.trainSet().size());
+        progressTracker.beginSubTask("Evaluate on train data");
+        progressTracker.setSteps(outerSplit.trainSet().size());
         registerMetricScores(outerSplit.trainSet(), bestClassifier, trainingStatistics::addOuterTrainScore, progressTracker);
         var outerTrainMetrics = trainingStatistics.winningModelOuterTrainMetrics();
         progressTracker.logMessage(formatWithLocale("Final model metrics on full train set: %s", outerTrainMetrics));
+        progressTracker.endSubTask("Evaluate on train data");
 
+        progressTracker.beginSubTask("Evaluate on test data");
+        progressTracker.setSteps(outerSplit.testSet().size());
         registerMetricScores(outerSplit.testSet(), bestClassifier, trainingStatistics::addTestScore, progressTracker);
         var testMetrics = trainingStatistics.winningModelTestMetrics();
         progressTracker.logMessage(formatWithLocale("Final model metrics on test set: %s", testMetrics));
-
         progressTracker.endSubTask("Evaluate on test data");
     }
 

pipeline/src/test/resources/expectedLogs/node-classification-log

@@ -33,12 +33,13 @@ MY DUMMY TASK :: Train best model :: Epoch 5 with loss 0.6128
 MY DUMMY TASK :: Train best model :: converged after 5 out of 100 epochs. Initial loss: 0.6931, Last loss: 0.6128.
 MY DUMMY TASK :: Train best model 100%
 MY DUMMY TASK :: Train best model :: Finished
+MY DUMMY TASK :: Evaluate on train data :: Start
+MY DUMMY TASK :: Evaluate on train data 100%
+MY DUMMY TASK :: Evaluate on train data :: Final model metrics on full train set: {F1_class_1=0.8235}
+MY DUMMY TASK :: Evaluate on train data :: Finished
 MY DUMMY TASK :: Evaluate on test data :: Start
-MY DUMMY TASK :: Evaluate on test data 50%
-MY DUMMY TASK :: Evaluate on test data :: Final model metrics on full train set: {F1_class_1=0.8235}
-MY DUMMY TASK :: Evaluate on test data 75%
-MY DUMMY TASK :: Evaluate on test data :: Final model metrics on test set: {F1_class_1=0.7499}
 MY DUMMY TASK :: Evaluate on test data 100%
+MY DUMMY TASK :: Evaluate on test data :: Final model metrics on test set: {F1_class_1=0.7499}
 MY DUMMY TASK :: Evaluate on test data :: Finished
 MY DUMMY TASK :: Retrain best model :: Start
 MY DUMMY TASK :: Retrain best model :: Initial loss 0.6931

pipeline/src/test/resources/expectedLogs/node-classification-with-range-log

@@ -33,12 +33,13 @@ MY DUMMY TASK :: Train best model :: Epoch 5 with loss 0.6128
 MY DUMMY TASK :: Train best model :: converged after 5 out of 100 epochs. Initial loss: 0.6931, Last loss: 0.6128.
 MY DUMMY TASK :: Train best model 100%
 MY DUMMY TASK :: Train best model :: Finished
+MY DUMMY TASK :: Evaluate on train data :: Start
+MY DUMMY TASK :: Evaluate on train data 100%
+MY DUMMY TASK :: Evaluate on train data :: Final model metrics on full train set: {F1_class_1=0.8235}
+MY DUMMY TASK :: Evaluate on train data :: Finished
 MY DUMMY TASK :: Evaluate on test data :: Start
-MY DUMMY TASK :: Evaluate on test data 50%
-MY DUMMY TASK :: Evaluate on test data :: Final model metrics on full train set: {F1_class_1=0.8235}
-MY DUMMY TASK :: Evaluate on test data 75%
-MY DUMMY TASK :: Evaluate on test data :: Final model metrics on test set: {F1_class_1=0.7499}
 MY DUMMY TASK :: Evaluate on test data 100%
+MY DUMMY TASK :: Evaluate on test data :: Final model metrics on test set: {F1_class_1=0.7499}
 MY DUMMY TASK :: Evaluate on test data :: Finished
 MY DUMMY TASK :: Retrain best model :: Start
 MY DUMMY TASK :: Retrain best model :: Initial loss 0.6931