如何在TensorFlow 2.0的Keras序列网络中使用LBFGS优化器？

Keras原生并未内置LBFGS优化器，需借助TensorFlow Probability（TFP）的tfp.lbfgs_minimize实现。LBFGS是全批量优化器，适合小数据集（会一次性使用全量数据计算损失和梯度），训练时无需分批处理。

实现步骤 #

1. 安装并导入依赖 #

确保安装tensorflow-probability，然后导入所需库：

import tensorflow as tf
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Dense
import tensorflow_probability as tfp

2. 定义模型与自定义损失 #

保持原模型结构，按需求实现自定义损失（示例为MSE结合参数正则项）：

# 自定义损失函数：MSE与其他变量的组合示例
def custom_loss(y_true, y_pred):
    mse_loss = tf.keras.losses.MSE(y_true, y_pred)
    # 替换为你的实际自定义逻辑，此处示例为参数L2正则项
    l2_reg = tf.add_n([tf.nn.l2_loss(var) for var in model.trainable_variables])
    return mse_loss + 0.01 * l2_reg

# 定义原模型结构
model = Sequential()
model.add(Dense(units=32, activation='relu', input_dim=1))
model.add(Dense(units=10, activation='softmax'))

3. 参数扁平化与辅助函数 #

LBFGS要求参数为扁平化向量，需实现以下辅助函数：

# 获取模型可训练变量
trainable_vars = model.trainable_variables
# 计算参数总长度
num_vars = sum([tf.size(v).numpy() for v in trainable_vars])

# 将扁平化向量恢复为模型参数形状的函数
def assign_params(params):
    idx = 0
    for var in trainable_vars:
        var_shape = var.shape
        var_size = tf.size(var).numpy()
        var.assign(tf.reshape(params[idx:idx+var_size], var_shape))
        idx += var_size

# 包装损失与梯度计算的函数，符合tfp.lbfgs_minimize的输入要求
def loss_and_grad_fn(params):
    assign_params(params)
    with tf.GradientTape() as tape:
        y_pred = model(x_train, training=True)
        loss = custom_loss(y_train, y_pred)
    grads = tape.gradient(loss, trainable_vars)
    # 将梯度转为扁平化向量
    flat_grads = tf.concat([tf.reshape(g, [-1]) for g in grads], axis=0)
    return loss, flat_grads

4. 执行LBFGS优化 #

初始化扁平化参数并调用优化器：

# 初始化扁平化参数
initial_params = tf.concat([tf.reshape(v, [-1]) for v in trainable_vars], axis=0)

# 执行LBFGS优化
results = tfp.optimizer.lbfgs_minimize(
    value_and_gradients_function=loss_and_grad_fn,
    initial_position=initial_params,
    num_iterations=100,  # 最大迭代次数，按需调整
    tolerance=1e-6       # 收敛阈值，损失变化小于该值时提前停止
)

# 将优化后的参数赋值回模型
assign_params(results.position)

5. 评估模型 #

优化完成后正常评估性能：

loss, accuracy = model.evaluate(x_test, y_test)
print('Test loss:', loss)
print('Test accuracy:', accuracy)

完整代码示例 #

import tensorflow as tf
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Dense
import tensorflow_probability as tfp

# 假设x_train, y_train, x_test, y_test已提前准备好

# 定义模型架构
model = Sequential()
model.add(Dense(units=32, activation='relu', input_dim=1))
model.add(Dense(units=10, activation='softmax'))

# 自定义损失函数：MSE与其他变量的组合
def custom_loss(y_true, y_pred):
    mse_loss = tf.keras.losses.MSE(y_true, y_pred)
    # 替换为你的实际自定义逻辑
    l2_reg = tf.add_n([tf.nn.l2_loss(var) for var in model.trainable_variables])
    return mse_loss + 0.01 * l2_reg

# 准备LBFGS优化辅助函数
trainable_vars = model.trainable_variables
num_vars = sum([tf.size(v).numpy() for v in trainable_vars])

def assign_params(params):
    idx = 0
    for var in trainable_vars:
        var_shape = var.shape
        var_size = tf.size(var).numpy()
        var.assign(tf.reshape(params[idx:idx+var_size], var_shape))
        idx += var_size

def loss_and_grad_fn(params):
    assign_params(params)
    with tf.GradientTape() as tape:
        y_pred = model(x_train, training=True)
        loss = custom_loss(y_train, y_pred)
    grads = tape.gradient(loss, trainable_vars)
    flat_grads = tf.concat([tf.reshape(g, [-1]) for g in grads], axis=0)
    return loss, flat_grads

# 执行LBFGS优化
initial_params = tf.concat([tf.reshape(v, [-1]) for v in trainable_vars], axis=0)
results = tfp.optimizer.lbfgs_minimize(
    value_and_gradients_function=loss_and_grad_fn,
    initial_position=initial_params,
    num_iterations=100,
    tolerance=1e-6
)

# 赋值优化后的参数
assign_params(results.position)

# 评估模型
loss, accuracy = model.evaluate(x_test, y_test)
print('Test loss:', loss)
print('Test accuracy:', accuracy)

注意事项 #

• LBFGS是全批量优化器，仅适合小数据集，大数据集使用会导致内存溢出。
• 调整num_iterations和tolerance控制优化过程：num_iterations为最大迭代次数，tolerance为损失收敛阈值。
• 自定义损失函数需基于TensorFlow原生操作，确保能被自动微分。

内容的提问来源于stack exchange，提问作者Ashwin Raikar