How to run scipy's BFGS on GPU

Question

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?

Asked By: jinus b

||

Source

Answer 1

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)

Answered By: Saman

How to run scipy's BFGS on GPU

Question:

Answers: