Skip to content
Advertisement

How to run scipy’s BFGS on GPU

I’d like to run scipy implementation of BFGS optimization algorithm on GPU and scipy seems not to support GPUs. The target function which I want to run on GPU is the following one which is part of the implementation of this repository:

//here the variable initializations

opts = {'gtol': effective_gamma, 'norm': 2}
result = minimize(private_loss, x0, (x, y), method='BFGS', jac=private_gradient, options=opts,
callback=cb)

I know there is Tensorflow Probablity implementation of BFGS, but I couldn’t find out how I can convert this scipy function into Tensordlow Probablity. Any Idea how I could to run the following function on GPU with minimum code change?

Advertisement

Answer

Here are two suggestions of mine:

  • jax.scipy: Jax contains the implementation of scipy and also supports GPUs and TPUs. You can theoretically install it, convert your numpy variables to jax.numpy and call jax.scipy.optimize.minimize(params):

     import jax.numpy as jnp
     from jax.scipy.optimize import minimize
    
     // here x0,x,y initialization and private_loss and private_gradient functions definition 
    
     x0 = jax.numpy.asarray(x0)
     x = jax.numpy.asarray(x)
     y = jax.numpy.asarray(y)
    
     opts = {'gtol': effective_gamma, 'norm': 2}
     result = minimize(private_loss, x0, (x, y), method='BFGS', jac=private_gradient, options=opts, callback=cb)  
    
     // here rest of the code
    

Dont foget to also take care of converting the variables which will be used in private_loss and private_gradient function to jax.numpy.

  • Tensorflow Probability: As you already mentioned, you can also use bfgs_minimize implementation by tensorflow. based on the colab here, your code will be something like this:

     def make_val_and_grad_fn(value_fn):
       @functools.wraps(value_fn)
       def val_and_grad(param):
         return tfp.math.value_and_gradient(value_fn, param)
       return val_and_grad
    
     def run(optimizer):
       result2 = optimizer()
       return np_value(result2)
    
     @make_val_and_grad_fn
     def private_loss(param):
         // here private_loss
    
     x_var = tf.Variable(x)
     y_var = tf.Variable(y, dtype = tf.float64)
     x0 = tf.Variable(x0)
    
     tolerance = effective_gamma
    
     @tf.function
     def minimize_with_bfgs():
         result1 = tfp.optimizer.bfgs_minimize(
         private_loss,
         initial_position=tf.constant(x0),
         tolerance=tolerance)
         return result1
    
     results = run(minimize_with_bfgs)
    
User contributions licensed under: CC BY-SA
10 People found this is helpful
Advertisement