How to run scipy's BFGS on GPU-CodePudding

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)