Question about running botorch on GPU

[A-ep93]

Hello all,
I have a botorch optimization code and I tried to modify my code to make sure that I am using GPU instead of CPU as much as possible.
Here are the main sections of my code:

tkwargs = {
    "dtype": torch.double,
    "device": torch.device("cuda" if torch.cuda.is_available() else "cpu"),
}


def problem(input):
        # This is the function I want to optimize
        return torch.tensor(result,**tkwargs)

bounds=torch.tensor([[0],[1]],**tkwargs)

def generate_initial_data(n=6):
        # generate training data
        train_x = draw_sobol_samples(
            bounds=bounds,n=1, q=n, seed=torch.randint(1000000, (1,)).item()
        ).squeeze(0)
        train_obj = problem(train_x)
        best=train_obj.max()
        return train_x, train_obj, best


def get_next_points(train_x, train_obj, best, bounds, n_points):
        models =SingleTaskGP(train_x,train_obj, outcome_transform=Standardize(m=1))
        mll = ExactMarginalLogLikelihood(models.likelihood, models)
        fit_gpytorch_model(mll)
        EI=ExpectedImprovement(
            model=models,
            best_f=best)
        candidate, _ =optimize_acqf(
            acq_function=EI,
            bounds=bounds,
            q=n_points,
            num_restarts=20,
            raw_samples=512,
            )
        return candidate


n_runs=50
train_x, train_obj, best=generate_initial_data(5)



for i in range(n_runs):
        new_candidate=get_next_points(train_x, train_obj, best, bounds, 1)
        new_results=problem(new_candidate)
        train_x=torch.cat([train_x,new_candidate])
        train_obj=torch.cat([train_obj,new_results])
        best=train_obj.max()
        location=train_x[torch.argmax(train_obj)]

Should I apply any other modifications to the code above to make sure all possible steps are running on GPU not CPU?
Any comments are greatly appreciated.

[saitcakmak]

As it's written, your code should be running on the device specified by tkwargs. So, if you have PyTorch setup with cuda support and you have an nvidia GPU available, it should be running on the GPU. You can check this by printing out tkwargs or torch.cuda.is_available().

If you're asking because you are not observing significant speed ups from switching to GPU, this is expected. You're using the analytic EI acquisition function, which does not use expensive matrix-vector multiplication, where the GPUs are typically much faster. We generally recommend GPUs for multi-objective optimization or for acquisition functions that require large q-batch sampling from the GP posterior.