Add structured readout versions of binary and multiclass classificati…

…on tasks. PiperOrigin-RevId: 588115018
tensorflow · Dec 5, 2023 · e5988bb · e5988bb
1 parent 68f1e28
commit e5988bb
Show file tree

Hide file tree

Showing 4 changed files with 467 additions and 26 deletions.
diff --git a/tensorflow_gnn/api_def/runner-symbols.txt b/tensorflow_gnn/api_def/runner-symbols.txt
@@ -26,6 +26,7 @@ runner.ParameterServerStrategy
 runner.PassthruDatasetProvider
 runner.PassthruSampleDatasetsProvider
 runner.Predictions
+runner.NodeBinaryClassification
 runner.RootNodeBinaryClassification
 runner.RootNodeLabelFn
 runner.RootNodeMeanAbsoluteError
@@ -34,6 +35,7 @@ runner.RootNodeMeanAbsolutePercentageError
 runner.RootNodeMeanSquaredError
 runner.RootNodeMeanSquaredLogScaledError
 runner.RootNodeMeanSquaredLogarithmicError
+runner.NodeMulticlassClassification
 runner.RootNodeMulticlassClassification
 runner.RunResult
 runner.SampleTFRecordDatasetsProvider

diff --git a/tensorflow_gnn/runner/__init__.py b/tensorflow_gnn/runner/__init__.py
@@ -94,10 +94,12 @@
 # in `distribute_test.py`.)
 #
 # Tasks (Classification)
-RootNodeBinaryClassification = classification.RootNodeBinaryClassification
-RootNodeMulticlassClassification = classification.RootNodeMulticlassClassification
 GraphBinaryClassification = classification.GraphBinaryClassification
 GraphMulticlassClassification = classification.GraphMulticlassClassification
+NodeBinaryClassification = classification.NodeBinaryClassification
+NodeMulticlassClassification = classification.NodeMulticlassClassification
+RootNodeBinaryClassification = classification.RootNodeBinaryClassification
+RootNodeMulticlassClassification = classification.RootNodeMulticlassClassification
 # Tasks (Link Prediction)
 DotProductLinkPrediction = link_prediction.DotProductLinkPrediction
 HadamardProductLinkPrediction = link_prediction.HadamardProductLinkPrediction