[DNR][outdated] Batch inference across estimators

benchmark/tabular/model.py

@@ -62,6 +62,7 @@ def _set_default_params(self) -> None:
         self._set_default_param_value("max_context_size", None)
         self._set_default_param_value("max_columns", None)
         self._set_default_param_value("kv_cache", False)
+        self._set_default_param_value("estimator_batch_size", "auto")
 
     def _fit(
         self,
@@ -127,6 +128,7 @@ def _fit(
             y_context = y_context[perm].unflatten(0, shape)
             num_estimators = None
         self._expand_query = num_estimators is None
+        self._context_shape = x_context.shape[-2:]
 
         recipe = self._create_recipe()
         if params["max_columns"] is not None:
@@ -146,6 +148,7 @@ def _fit(
                     y=y_context,
                     recipe=recipe,
                     num_estimators=num_estimators,
+                    estimator_batch_size=1,
                     generator=generator,
                 )
             return
@@ -202,7 +205,10 @@ def _predict_proba(
                 self.autocast_dtype,
                 enabled=x_query.is_cuda,
             ):
-                out = self.model.predict(x_query)
+                out = self.model.predict(
+                    x_query,
+                    estimator_batch_size=self._estimator_batch_size(x_query),
+                )
         else:
             with (
                 torch.inference_mode(),
@@ -215,24 +221,20 @@ def _predict_proba(
                 generator = torch.Generator(self._device).set_state(
                     self._rng_state
                 )
-                outputs = []
-                for context, query in zip(self._contexts, queries):
-                    with torch.amp.autocast(
-                        self._device.type,
-                        self.autocast_dtype,
-                        enabled=x_query.is_cuda,
-                    ):
-                        out = self.model._forward(
-                            x_context=context.x.to(self._device),
-                            y_context=context.y.to(self._device),
-                            x_query=query.x,
-                            related_context_tables=None,
-                            related_query_tables=None,
-                            cache=None,
-                            generator=generator,
-                            _schema=self._schema,
-                        )
-                    outputs.append(out.to(query.x.dtype))
+                with torch.amp.autocast(
+                    self._device.type,
+                    self.autocast_dtype,
+                    enabled=x_query.is_cuda,
+                ):
+                    outputs = self.model._forward_estimators(
+                        contexts=self._contexts,
+                        queries=queries,
+                        estimator_batch_size=self._estimator_batch_size(
+                            x_query
+                        ),
+                        generator=generator,
+                        _schema=self._schema,
+                    )
 
                 if self.problem_type == REGRESSION:
                     outputs = list(
@@ -251,6 +253,28 @@ def _predict_proba(
         probabilities = out.numerical[..., indices].float().cpu().numpy()
         return self._convert_proba_to_unified_form(probabilities)
 
+    def _estimator_batch_size(self, x: torch.Tensor) -> int | None:
+        params = self._get_model_params()
+        estimator_batch_size = params["estimator_batch_size"]
+        if estimator_batch_size != "auto":
+            return estimator_batch_size
+        # Subsampled contexts can produce different cache shapes per estimator.
+        if not x.is_cuda or self._expand_query:
+            return 1
+
+        num_rows, num_cols = self._context_shape
+        if not params["kv_cache"]:
+            # Uncached inference processes context and query rows together.
+            num_rows += x.size(-2)
+            if num_rows > 3_000 or num_rows * num_cols > 50_000:
+                return 1
+            return self._num_estimators
+
+        num_rows = max(num_rows, x.size(-2))
+        if num_rows > 2_000 or num_rows * num_cols >= 50_000:
+            return 1
+        return self._num_estimators
+
     def get_device(self) -> str:
         return str(next(self.model.parameters()).device)