Home > Mobile >  How to run scipy's BFGS on GPU
How to run scipy's BFGS on GPU

Time:02-13

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?

CodePudding user response:

Here are two suggestions of mine:

  • jax.scipy: Jax contains the implementation of scipy and also supports GPUs and TPUs. You can theoricatically 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)
    
  • Related