与 TensorFlow 并行执行许多基于梯度的优化

Question

我有一个模型需要用 tfp.math.ode.BDF 求解 ODE 系统，我想找到该模型对 n > 1000 个数据集的单个最小二乘拟合。也就是说，如果我的模型有 m 个参数，那么在优化过程结束时，我将有一个 n x m 张量的最佳拟合参数值。

并行执行此优化的最佳方式是什么？此时我打算定义一个 objective 函数，将 n 个单独的残差平方和相加，然后使用 tfp.optimizer.lbfgs_minimize 找到组合的 n×m 参数的最佳拟合值。

Answer 1

这不会很有帮助，我不知道 ODE 是什么，但是如果我试图针对 n 个模型（无论是线性回归还是深度学习或任何其他通过梯度优化的模型），我将简单地实例化模型的 n 个副本，将所有 n 个数据集前向支持对应的模型，并关键地总结损失。自动微分会处理剩下的事情。

除了最后的损失求和之外，n 个模型不会相互影响，因此自动差异也会使梯度保持隔离。我有一段时间没有使用 l-bfgs，但 IIRC 它就像 SGD，除了它在二阶偏导数中的因素，所以它可能是用渐变带内的渐变带实现的。

这是构建包含 n=2 个模型的单个张量流模型的最小示例。每个模型包含 m=2 个权重（即 y=mx+b 方程中的 m 和 b）。您训练 n 个数据集，这些数据集与 n 个模型具有一对一的关系。你将从中得到 n x m 个权重。

import tensorflow as tf


def get_model():
  """Create linear regression model of the form y=mx+b."""
  model = tf.keras.Sequential()
  model.add(tf.keras.layers.Dense(1, input_shape=(1,)))
  return model

# Setup two different datasets. 
x1 = tf.range(start=-1., limit=1., delta=0.001)
y1 = x1 * .1 + .2

x2 = tf.range(start=-1., limit=1., delta=0.001)
y2 = x2 * .3 + .4


def create_parallel_model():
  """Creates a single model that hold multiple independent models.
  
  Returns:
    parallel_model: The single parallel model holding multiple models. 
    model1: First componenent model. 
    model2: Second componenent model. 
    """
  model1 = get_model()
  model2 = get_model()

  inp1 = tf.keras.Input((1,), name='inp1')
  inp2 = tf.keras.Input((1,), name='inp2')

  outp1 = model1(inp1)
  outp2 = model2(inp2)

  parallel_model = tf.keras.Model([inp1, inp2], [outp1, outp2])

  return parallel_model, model1, model2


pmodel, model1, model2 = create_parallel_model()
pmodel.compile('sgd', 'mse')
pmodel.fit(x=[x1, x2], y=[y1, y2], epochs=100)

print("First model's dense layer, m and b should be 0.1 and 0.2.")
print(model1.layers[0].get_weights())

print("Second model's dense layer, m and b should be 0.3 and 0.4.")
print(model2.layers[0].get_weights())

您的输出应该如下所示

First model's dense layer, m and b should be 0.1 and 0.2.
[array([[0.10000037]], dtype=float32), array([0.20000002], dtype=float32)]
Second model's dense layer, m and b should be 0.3 and 0.4.
[array([[0.30000147]], dtype=float32), array([0.39999998], dtype=float32)]

另一种方法是在单个进程中一次简单地训练一个模型和数据集，但运行并行训练多个进程。如果你使用 TF 的动态内存分配，TF 不会占用所有的 GPU 内存并且 Nvidia GPU 支持多处理 (https://docs.nvidia.com/deploy/mps/index.html).

physical_devices = tf.config.list_physical_devices('GPU') 
for gpu_instance in physical_devices: 
    tf.config.experimental.set_memory_growth(gpu_instance, True)

Answer 2

我相信 BDF 优化器和 LBFGS 都支持（问题的）批处理，因此您可以有一个外部“问题”轴到您的数据和 leastsq return价值。但是由于 BDF 是针对刚性问题的，因此某些问题的运行时间可能比其他问题长得多，您最好分别处理每个问题 (tf.map_fn) 而不是 [=18] =] 将它们全部并行 - 在批处理中，在计算问题 Y 的 BDF 积分之前，您不能运行提前进入问题 X 的下一个 LBFGS 迭代。或者只使用 python for 遍历你的问题，每次调用 @tf.function def lbfgs_over_bdf(data): ....

与 TensorFlow 并行执行许多基于梯度的优化

performing many gradient-based optimizations in parallel with TensorFlow

python

tensorflow

tensorflow-probability