优化器

优化器#

MLX 中的优化器既可以与 mlx.nn 一起使用，也可以与纯粹的 mlx.core 函数一起使用。一个典型的例子是调用 Optimizer.update() 根据损失梯度更新模型的参数，然后调用 mlx.core.eval() 求值模型的参数和**优化器状态**。

# Create a model
model = MLP(num_layers, train_images.shape[-1], hidden_dim, num_classes)
mx.eval(model.parameters())

# Create the gradient function and the optimizer
loss_and_grad_fn = nn.value_and_grad(model, loss_fn)
optimizer = optim.SGD(learning_rate=learning_rate)

for e in range(num_epochs):
    for X, y in batch_iterate(batch_size, train_images, train_labels):
        loss, grads = loss_and_grad_fn(model, X, y)

        # Update the model with the gradients. So far no computation has happened.
        optimizer.update(model, grads)

        # Compute the new parameters but also the optimizer state.
        mx.eval(model.parameters(), optimizer.state)

保存和加载#

要序列化优化器，请保存其状态。要加载优化器，请加载并设置保存的状态。这里有一个简单的例子

import mlx.core as mx
from mlx.utils import tree_flatten, tree_unflatten
import mlx.optimizers as optim

optimizer = optim.Adam(learning_rate=1e-2)

# Perform some updates with the optimizer
model = {"w" : mx.zeros((5, 5))}
grads = {"w" : mx.ones((5, 5))}
optimizer.update(model, grads)

# Save the state
state = tree_flatten(optimizer.state)
mx.save_safetensors("optimizer.safetensors", dict(state))

# Later on, for example when loading from a checkpoint,
# recreate the optimizer and load the state
optimizer = optim.Adam(learning_rate=1e-2)

state = tree_unflatten(list(mx.load("optimizer.safetensors").items()))
optimizer.state = state

请注意，并非所有优化器配置参数都会保存在状态中。例如，对于 Adam，学习率会保存，但 betas 和 eps 参数则不会。一个好的经验法则是，如果参数可以被调度，那么它就会包含在优化器状态中。

clip_grad_norm(grads, max_norm)

裁剪梯度的全局范数。
目录