优化器

优化器#

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)

裁剪梯度的全局范数。
Contents