如何在 GPU 上 运行 scipy 的 BFGS
How to run scipy's BFGS on GPU
我想 运行 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 上 运行 以下功能吗？
这里有两个我的建议：
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)
我想 运行 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 上 运行 以下功能吗？
这里有两个我的建议：
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)