如何在 GPU 上运行 scipy 的 BFGS

Question

我想运行 scipy 在 GPU 上实现 BFGS 优化算法，scipy 似乎不支持 GPU。我想在 GPU 上运行的目标函数是以下函数，它是 this 存储库实现的一部分：

//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)

我知道有 BFGS 的 Tensorflow Probablity 实现，但我找不到如何将此 scipy 函数转换为 Tensordlow Probablity。知道如何以最少的代码更改在 GPU 上运行以下功能吗？

Answer 1

这里有两个我的建议：

jax.scipy: Jax 包含 scipy 的实现，还支持 GPU 和 TPU。您可以理论上安装它，将您的 numpy 变量转换为 jax.numpy 并调用 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

不要忘记将 private_loss 和 private_gradient 函数中使用的变量转换为 jax.numpy。

Tensorflow Probability：正如您已经提到的，您也可以使用 bfgs_minimize 通过 tensorflow 实现。基于此处的 colab，您的代码将如下所示：

 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)

如何在 GPU 上运行 scipy 的 BFGS

How to run scipy's BFGS on GPU

python

tensorflow

scipy

tensorflow-probability

如何在 GPU 上 运行 scipy 的 BFGS

How to run scipy's BFGS on GPU

python

tensorflow

scipy

tensorflow-probability

如何在 GPU 上运行 scipy 的 BFGS