[AIR] Introduce better scoring API for BatchPredictor (#26451)

Signed-off-by: Amog Kamsetty <[email protected]>

As discussed offline, allow configurability for feature columns and keep columns in BatchPredictor for better scoring UX on test datasets.

doc/source/ray-air/examples/torch_image_example.ipynb

@@ -435,7 +435,7 @@
     ")\n",
     "\n",
     "outputs: ray.data.Dataset = batch_predictor.predict(\n",
-    "    data=predict_dataset, dtype=torch.float\n",
+    "    data=test_dataset, dtype=torch.float, feature_columns=[\"image\"], keep_columns=[\"label\"]\n",
     ")"
    ]
   },
@@ -482,7 +482,7 @@
     "def convert_logits_to_classes(df):\n",
     "    best_class = df[\"predictions\"].map(lambda x: x.argmax())\n",
     "    df[\"prediction\"] = best_class\n",
-    "    return df[[\"prediction\"]]\n",
+    "    return df\n",
     "\n",
     "predictions = outputs.map_batches(\n",
     "    convert_logits_to_classes, batch_format=\"pandas\"\n",
@@ -536,7 +536,7 @@
     "    df[\"correct\"] = df[\"prediction\"] == df[\"label\"]\n",
     "    return df[[\"prediction\", \"label\", \"correct\"]]\n",
     "\n",
-    "scores = test_dataset.zip(predictions).map_batches(calculate_prediction_scores)\n",
+    "scores = predictions.map_batches(calculate_prediction_scores)\n",
     "\n",
     "scores.show(1)"
    ]

doc/source/ray-air/examples/torch_incremental_learning.ipynb

@@ -640,7 +640,7 @@
     "\n",
     "  batch_predictor = BatchPredictor.from_checkpoint(checkpoint, predictor_cls=TorchPredictor, model=SimpleMLP(num_classes=10))\n",
     "  model_output = batch_predictor.predict(\n",
-    "            data=test_dataset, feature_columns=[\"image\"], unsqueeze=False\n",
+    "            data=test_dataset, feature_columns=[\"image\"], keep_columns=[\"label\"]\n",
     "        )\n",
     "  \n",
     "  # Postprocess model outputs.\n",
@@ -654,12 +654,10 @@
     "  \n",
     "  # Then, for each prediction output, see if it matches with the ground truth\n",
     "  # label.\n",
-    "  zipped_dataset = test_dataset.zip(prediction_results)\n",
-    "\n",
     "  def calculate_prediction_scores(df):\n",
     "      return pd.DataFrame({\"correct\": df[\"predictions\"] == df[\"label\"]})\n",
     "\n",
-    "  correct_dataset = zipped_dataset.map_batches(\n",
+    "  correct_dataset = prediction_results.map_batches(\n",
     "      calculate_prediction_scores, batch_format=\"pandas\"\n",
     "  )\n",
     "\n",

python/ray/experimental/state/api.py

@@ -1,10 +1,10 @@
+import logging
+import threading
 import urllib
 import warnings
-import threading
-import logging
-from dataclasses import fields
-from typing import Dict, Generator, List, Optional, Tuple, Union, Any
 from contextlib import contextmanager
+from dataclasses import fields
+from typing import Any, Dict, Generator, List, Optional, Tuple, Union
 
 import requests
 

python/ray/train/batch_predictor.py

@@ -1,5 +1,5 @@
 import inspect
-from typing import Any, Dict, Optional, Type, Union
+from typing import Any, Dict, Optional, List, Type, Union
 
 import ray
 from ray.air import Checkpoint
@@ -45,6 +45,8 @@ def predict(
         self,
         data: Union[ray.data.Dataset, ray.data.DatasetPipeline],
         *,
+        feature_columns: Optional[List[str]] = None,
+        keep_columns: Optional[List[str]] = None,
         batch_size: int = 4096,
         min_scoring_workers: int = 1,
         max_scoring_workers: Optional[int] = None,
@@ -61,24 +63,41 @@ def predict(
             >>> from ray.air import Checkpoint
             >>> from ray.train.predictor import Predictor
             >>> from ray.train.batch_predictor import BatchPredictor
-            >>> # Create a dummy predictor that always returns `42` for each input.
+            >>> # Create a dummy predictor that returns identity as the predictions.
             >>> class DummyPredictor(Predictor):
             ...     @classmethod
             ...     def from_checkpoint(cls, checkpoint, **kwargs):
             ...         return cls()
-            ...     def predict(self, data, **kwargs):
-            ...         return pd.DataFrame({"a": [42] * len(data)})
+            ...     def _predict_pandas(self, data_df, **kwargs):
+            ...         return data_df
             >>> # Create a batch predictor for this dummy predictor.
             >>> batch_pred = BatchPredictor( # doctest: +SKIP
             ...     Checkpoint.from_dict({"x": 0}), DummyPredictor)
             >>> # Create a dummy dataset.
-            >>> ds = ray.data.range_tensor(1000, parallelism=4) # doctest: +SKIP
+            >>> ds = ray.data.from_pandas(pd.DataFrame({ # doctest: +SKIP
+            ...     "feature_1": [1, 2, 3], "label": [1, 2, 3]}))
             >>> # Execute batch prediction using this predictor.
-            >>> print(batch_pred.predict(ds)) # doctest: +SKIP
-            Dataset(num_blocks=4, num_rows=1000, schema={a: int64})
+            >>> predictions = batch_pred.predict(ds, # doctest: +SKIP
+            ...     feature_columns=["feature_1"], keep_columns=["label"])
+            >>> print(predictions) # doctest: +SKIP
+            Dataset(num_blocks=1, num_rows=3, schema={a: int64, label: int64})
+            >>> # Calculate final accuracy.
+            >>> def calculate_accuracy(df):
+            ...    return pd.DataFrame({"correct": df["predictions"] == df["label"]})
+            >>> correct = predictions.map_batches(calculate_accuracy) # doctest: +SKIP
+            >>> print("Final accuracy: ", # doctest: +SKIP
+            ...    correct.sum(on="correct") / correct.count())
+            Final accuracy: 1.0000
 
         Args:
             data: Ray dataset or pipeline to run batch prediction on.
+            feature_columns: List of columns in data to use for prediction. Columns not
+                specified will be dropped from `data` before being passed to the
+                predictor. If None, use all columns.
+            keep_columns: List of columns in `data` to include in the prediction result.
+                This is useful for calculating final accuracies/metrics on the result
+                dataset. If None, the columns in the output dataset will contain just
+                the prediction results.
             batch_size: Split dataset into batches of this size for prediction.
             min_scoring_workers: Minimum number of scoring actors.
             max_scoring_workers: If set, specify the maximum number of scoring actors.
@@ -112,7 +131,15 @@ def __init__(self):
                 )
 
             def __call__(self, batch):
-                prediction_output = self.predictor.predict(batch, **predict_kwargs)
+                if feature_columns:
+                    prediction_batch = batch[feature_columns]
+                else:
+                    prediction_batch = batch
+                prediction_output = self.predictor.predict(
+                    prediction_batch, **predict_kwargs
+                )
+                if keep_columns:
+                    prediction_output[keep_columns] = batch[keep_columns]
                 return convert_batch_type_to_pandas(prediction_output)
 
         compute = ray.data.ActorPoolStrategy(
@@ -123,21 +150,25 @@ def __call__(self, batch):
         ray_remote_args["num_cpus"] = num_cpus_per_worker
         ray_remote_args["num_gpus"] = num_gpus_per_worker
 
-        return data.map_batches(
+        prediction_results = data.map_batches(
             ScoringWrapper,
             compute=compute,
             batch_format="pandas",
             batch_size=batch_size,
             **ray_remote_args,
         )
 
+        return prediction_results
+
     def predict_pipelined(
         self,
         data: ray.data.Dataset,
         *,
         blocks_per_window: Optional[int] = None,
         bytes_per_window: Optional[int] = None,
         # The remaining args are from predict().
+        feature_columns: Optional[List[str]] = None,
+        keep_columns: Optional[List[str]] = None,
         batch_size: int = 4096,
         min_scoring_workers: int = 1,
         max_scoring_workers: Optional[int] = None,
@@ -188,6 +219,13 @@ def predict_pipelined(
                 This will be treated as an upper bound for the window size, but each
                 window will still include at least one block. This is mutually
                 exclusive with ``blocks_per_window``.
+            feature_columns: List of columns in data to use for prediction. Columns not
+                specified will be dropped from `data` before being passed to the
+                predictor. If None, use all columns.
+            keep_columns: List of columns in `data` to include in the prediction result.
+                This is useful for calculating final accuracies/metrics on the result
+                dataset. If None, the columns in the output dataset will contain just
+                the prediction results.
             batch_size: Split dataset into batches of this size for prediction.
             min_scoring_workers: Minimum number of scoring actors.
             max_scoring_workers: If set, specify the maximum number of scoring actors.
@@ -215,6 +253,8 @@ def predict_pipelined(
         return self.predict(
             pipe,
             batch_size=batch_size,
+            feature_columns=feature_columns,
+            keep_columns=keep_columns,
             min_scoring_workers=min_scoring_workers,
             max_scoring_workers=max_scoring_workers,
             num_cpus_per_worker=num_cpus_per_worker,

python/ray/train/tests/test_batch_predictor.py

@@ -96,6 +96,37 @@ def test_batch_prediction_fs():
     )
 
 
+def test_batch_prediction_feature_cols():
+    batch_predictor = BatchPredictor.from_checkpoint(
+        Checkpoint.from_dict({"factor": 2.0}), DummyPredictor
+    )
+
+    test_dataset = ray.data.from_pandas(pd.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6]}))
+
+    assert batch_predictor.predict(
+        test_dataset, feature_columns=["a"]
+    ).to_pandas().to_numpy().squeeze().tolist() == [4.0, 8.0, 12.0]
+
+
+def test_batch_prediction_keep_cols():
+    batch_predictor = BatchPredictor.from_checkpoint(
+        Checkpoint.from_dict({"factor": 2.0}), DummyPredictor
+    )
+
+    test_dataset = ray.data.from_pandas(
+        pd.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6], "c": [7, 8, 9]})
+    )
+
+    output_df = batch_predictor.predict(
+        test_dataset, feature_columns=["a"], keep_columns=["b"]
+    ).to_pandas()
+
+    assert set(output_df.columns) == {"a", "b"}
+
+    assert output_df["a"].tolist() == [4.0, 8.0, 12.0]
+    assert output_df["b"].tolist() == [4, 5, 6]
+
+
 def test_automatic_enable_gpu_from_num_gpus_per_worker():
     """
     Test we automatically set underlying Predictor creation use_gpu to True if