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)