GPU support for generating points

Summary:
Strategy can now move the model to the gpu to allow generators to enjoy speedups.

This behavior is not owned by the generator, but by the strategy. Future iterations can move this.

Differential Revision: D65361838

aepsych/acquisition/bvn.py

@@ -72,6 +72,7 @@ def _bvnu(
     hk = h * k
 
     x, w = _gauss_legendre20(dtype=dh.dtype)
+    x, w = x.to(dh), w.to(dh)
 
     asr = 0.5 * torch.asin(r)
     sn = torch.sin(asr[..., None] * x)

aepsych/acquisition/lookahead.py

@@ -34,7 +34,7 @@ def Hb(p: Tensor) -> Tensor:
 
     Returns: Binary entropy for each probability.
     """
-    epsilon = torch.tensor(np.finfo(float).eps)
+    epsilon = torch.tensor(np.finfo(float).eps).to(p)
     p = torch.clamp(p, min=epsilon, max=1 - epsilon)
     return -torch.nan_to_num(p * torch.log2(p) + (1 - p) * torch.log2(1 - p))
 
@@ -78,6 +78,8 @@ def SUR_fn(Px: Tensor, P1: Tensor, P0: Tensor, py1: Tensor) -> Tensor:
 
     Returns: (b) tensor of SUR values.
     """
+    P1 = P1.to(Px)
+    py1 = py1.to(Px)
     sur = ClassErr(Px) - py1 * ClassErr(P1) - (1 - py1) * ClassErr(P0)
     return sur.sum(dim=-1)
 

aepsych/acquisition/lookahead_utils.py

@@ -144,7 +144,7 @@ def lookahead_inner(f_q: Normal) -> Tensor:
     pstar_marginal_0 = 1 - pstar_marginal_1
     pq_marginal_1 = probit(Mu_q / torch.sqrt(1 + Sigma2_q))
 
-    quad = GaussHermiteQuadrature1D()
+    quad = GaussHermiteQuadrature1D().to(Mu_q)
     fq_mvn = Normal(Mu_q, torch.sqrt(Sigma2_q))
     joint_ystar1_yq1 = quad(lookahead_inner, fq_mvn)
     joint_ystar0_yq1 = pq_marginal_1 - joint_ystar1_yq1

aepsych/strategy.py

@@ -26,7 +26,12 @@
 
 import numpy as np
 import torch
-
+from aepsych.acquisition import (
+    MonotonicBernoulliMCMutualInformation,
+    MonotonicMCLSE,
+    MonotonicMCPosteriorVariance,
+)
+from aepsych.acquisition.monotonic_rejection import MonotonicMCAcquisition
 from aepsych.config import Config
 from aepsych.generators.base import AEPsychGenerator
 from aepsych.generators.sobol_generator import SobolGenerator
@@ -62,6 +67,13 @@ class Strategy(object):
 
     _n_eval_points: int = 1000
 
+    no_gpu_acqfs = (
+        MonotonicMCAcquisition,
+        MonotonicBernoulliMCMutualInformation,
+        MonotonicMCPosteriorVariance,
+        MonotonicMCLSE,
+    )
+
     def __init__(
         self,
         generator: AEPsychGenerator,
@@ -74,6 +86,7 @@ def __init__(
         min_asks: int = 0,
         model: Optional[AEPsychMixin] = None,
         model_gpu: bool = False,
+        generator_gpu: bool = False,
         refit_every: int = 1,
         min_total_outcome_occurrences: int = 1,
         max_asks: Optional[int] = None,
@@ -94,6 +107,7 @@ def __init__(
             min_asks (int): The minimum number of points that should be generated from this strategy.
             model (ModelProtocol, optional): The AEPsych model of the data.
             model_gpu (bool): Whether to move the model to GPU, defaults to False.
+            generator_gpu (bool): Whether to use the GPU for generating points, defaults to False.
             refit_every (int): How often to refit the model from scratch.
             min_total_outcome_occurrences (int): The minimum number of total observations needed for each outcome before the strategy will finish.
                 Defaults to 1 (i.e., for binary outcomes, there must be at least one "yes" trial and one "no" trial).
@@ -143,6 +157,27 @@ def __init__(
 
             self.model_device = torch.device("cuda" if model_gpu else "cpu")
 
+        if generator_gpu:
+            if model is None:
+                logger.warning(
+                    f"GPU requested for generator {type(generator).__name__} but this generator has no model to move to GPU."
+                )
+                self.generator_device = torch.device("cpu")
+            else:
+                assert (
+                    torch.cuda.is_available()
+                ), f"GPU requested for generator {type(generator).__name__} but GPU is not found!"
+
+                if hasattr(generator, "acqf") and isinstance(
+                    generator.acqf, self.no_gpu_acqfs
+                ):
+                    logger.warning(f"{generator.acqf.__name__} does not support GPU")
+                    self.generator_device = torch.device("cpu")
+                else:
+                    self.generator_device = torch.device("cuda")
+        else:
+            self.generator_device = torch.device("cpu")
+
         self.run_indefinitely = run_indefinitely
         self.lb, self.ub, self.dim = _process_bounds(lb, ub, dim)
         self.min_total_outcome_occurrences = min_total_outcome_occurrences
@@ -237,8 +272,18 @@ def gen(self, num_points: int = 1) -> torch.Tensor:
         Returns:
             torch.Tensor: Next set of point(s) to evaluate, [num_points x dim].
         """
+        original_device = None
+        if self.model is not None and self.generator_device.type == "cuda":
+            original_device = self.model.device
+            self.model.to(self.generator_device)  # type: ignore
+
         self._count = self._count + num_points
-        return self.generator.gen(num_points, self.model)
+        points = self.generator.gen(num_points, self.model)
+
+        if original_device is not None:
+            self.model.to(original_device)  # type: ignore
+
+        return points
 
     @ensure_model_is_fresh
     def get_max(
@@ -452,6 +497,7 @@ def from_config(cls, config: Config, name: str) -> Strategy:
 
         gen_cls = config.getobj(name, "generator", fallback=SobolGenerator)
         generator = gen_cls.from_config(config)
+        generator_gpu = config.getboolean(gen_cls.__name__, "use_gpu", fallback=False)
 
         model_cls = config.getobj(name, "model", fallback=None)
         if model_cls is not None:
@@ -507,6 +553,7 @@ def from_config(cls, config: Config, name: str) -> Strategy:
             dim=dim,
             model=model,
             model_gpu=model_gpu,
+            generator_gpu=generator_gpu,
             generator=generator,
             min_asks=min_asks,
             refit_every=refit_every,

aepsych/utils.py

@@ -22,7 +22,7 @@ def make_scaled_sobol(
     lb: torch.Tensor, ub: torch.Tensor, size: int, seed: Optional[int] = None
 ) -> torch.Tensor:
     lb, ub, ndim = _process_bounds(lb, ub, None)
-    grid = SobolEngine(dimension=ndim, scramble=True, seed=seed).draw(size)
+    grid = SobolEngine(dimension=ndim, scramble=True, seed=seed).draw(size).to(lb)
 
     # rescale from [0,1] to [lb, ub]
     grid = lb + (ub - lb) * grid
@@ -127,7 +127,6 @@ def interpolate_monotonic(x, y, z, min_x=-np.inf, max_x=np.inf):
     y1 = y[idx]
 
     x_star = x0 + (x1 - x0) * (z - y0) / (y1 - y0)
-
     return x_star
 
 

tests/generators/test_optimize_acqf_generator.py

@@ -10,16 +10,45 @@
 
 import numpy as np
 import torch
-from aepsych.acquisition import MCLevelSetEstimation
+from aepsych.acquisition import (
+    ApproxGlobalSUR,
+    EAVC,
+    GlobalMI,
+    GlobalSUR,
+    LocalMI,
+    LocalSUR,
+    MCLevelSetEstimation,
+    MCPosteriorVariance,
+)
+from aepsych.acquisition.lookahead import MOCU, SMOCU
+from aepsych.acquisition.mutual_information import BernoulliMCMutualInformation
 from aepsych.config import Config
 from aepsych.generators import OptimizeAcqfGenerator
-from aepsych.models import (
-    GPClassificationModel,
-    PairwiseProbitModel,
-)
+from aepsych.models import GPClassificationModel, PairwiseProbitModel
+from aepsych.strategy import Strategy
 from botorch.acquisition.preference import AnalyticExpectedUtilityOfBestOption
+from parameterized import parameterized
 from sklearn.datasets import make_classification
 
+acqf_kwargs_target = {"target": 0.75}
+acqf_kwargs_lookahead = {"target": 0.75, "lookahead_type": "posterior"}
+
+acqfs = [
+    (MCPosteriorVariance, {}),
+    (ApproxGlobalSUR, acqf_kwargs_target),
+    (MOCU, acqf_kwargs_target),
+    (SMOCU, acqf_kwargs_target),
+    (EAVC, acqf_kwargs_target),
+    (EAVC, acqf_kwargs_lookahead),
+    (GlobalMI, acqf_kwargs_target),
+    (GlobalMI, acqf_kwargs_lookahead),
+    (GlobalSUR, acqf_kwargs_target),
+    (LocalMI, acqf_kwargs_target),
+    (LocalSUR, acqf_kwargs_target),
+    (MCLevelSetEstimation, acqf_kwargs_target),
+    (BernoulliMCMutualInformation, {}),
+]
+
 
 class TestOptimizeAcqfGenerator(unittest.TestCase):
     def test_time_limits(self):
@@ -84,6 +113,33 @@ def test_instantiate_eubo(self):
         acqf = generator._instantiate_acquisition_fn(model=model)
         self.assertTrue(isinstance(acqf, AnalyticExpectedUtilityOfBestOption))
 
+    @unittest.skipUnless(torch.cuda.is_available(), "no gpu available")
+    @parameterized.expand(acqfs)
+    def test_gpu_smoketest(self, acqf, acqf_kwargs):
+        lb = torch.tensor([0])
+        ub = torch.tensor([1])
+        model = GPClassificationModel(
+            lb=lb, ub=ub, inducing_size=10, inducing_point_method="pivoted_chol"
+        )
+
+        generator = OptimizeAcqfGenerator(acqf=acqf, acqf_kwargs=acqf_kwargs)
+
+        strat = Strategy(
+            lb=torch.tensor([0]),
+            ub=torch.tensor([1]),
+            model=model,
+            generator=generator,
+            stimuli_per_trial=1,
+            outcome_types=["binary"],
+            min_asks=1,
+            model_gpu=True,
+            generator_gpu=True,
+        )
+
+        strat.add_data(x=torch.tensor([0.90]), y=torch.tensor([1.0]))
+
+        strat.gen(1)
+
 
 if __name__ == "__main__":
     unittest.main()