多层感知机

多层感知机#

在此示例中，我们将通过实现一个简单的多层感知机来对 MNIST 进行分类，从而学习使用 mlx.nn。

第一步是导入所需的 MLX 包

import mlx.core as mx
import mlx.nn as nn
import mlx.optimizers as optim

import numpy as np

模型被定义为 MLP 类，它继承自 mlx.nn.Module。我们遵循标准惯例来创建一个新模块

  1. 定义 __init__ 方法，用于设置参数和/或子模块。有关 mlx.nn.Module 如何注册参数的更多信息，请参见Module 类文档。

  2. 定义 __call__ 方法，用于实现计算逻辑。

class MLP(nn.Module):
    def __init__(
        self, num_layers: int, input_dim: int, hidden_dim: int, output_dim: int
    ):
        super().__init__()
        layer_sizes = [input_dim] + [hidden_dim] * num_layers + [output_dim]
        self.layers = [
            nn.Linear(idim, odim)
            for idim, odim in zip(layer_sizes[:-1], layer_sizes[1:])
        ]

    def __call__(self, x):
        for l in self.layers[:-1]:
            x = mx.maximum(l(x), 0.0)
        return self.layers[-1](x)

我们定义损失函数，它计算每个样本的交叉熵损失的平均值。mlx.nn.losses 子包提供了一些常用损失函数的实现。

def loss_fn(model, X, y):
    return mx.mean(nn.losses.cross_entropy(model(X), y))

我们还需要一个函数来计算模型在验证集上的准确率

def eval_fn(model, X, y):
    return mx.mean(mx.argmax(model(X), axis=1) == y)

接下来，设置问题参数并加载数据。要加载数据，您需要我们的 mnist 数据加载器，我们将它导入为 mnist。

num_layers = 2
hidden_dim = 32
num_classes = 10
batch_size = 256
num_epochs = 10
learning_rate = 1e-1

# Load the data
import mnist
train_images, train_labels, test_images, test_labels = map(
    mx.array, mnist.mnist()
)

由于我们使用 SGD，我们需要一个迭代器，用于打乱训练集中的样本并构建小批量数据。

def batch_iterate(batch_size, X, y):
    perm = mx.array(np.random.permutation(y.size))
    for s in range(0, y.size, batch_size):
        ids = perm[s : s + batch_size]
        yield X[ids], y[ids]

最后，我们将模型、mlx.optimizers.SGD 优化器实例化，然后运行训练循环。

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

# Get a function which gives the loss and gradient of the
# loss with respect to the model's trainable parameters
loss_and_grad_fn = nn.value_and_grad(model, loss_fn)

# Instantiate the optimizer
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 optimizer state and model parameters
        # in a single call
        optimizer.update(model, grads)

        # Force a graph evaluation
        mx.eval(model.parameters(), optimizer.state)

    accuracy = eval_fn(model, test_images, test_labels)
    print(f"Epoch {e}: Test accuracy {accuracy.item():.3f}")

注意

mlx.nn.value_and_grad() 函数是一个便捷函数，用于获取损失相对于模型可训练参数的梯度。请勿将其与 mlx.core.value_and_grad() 混淆。

模型在训练集上仅训练几个周期后，应能达到不错的准确率（约 95%）。完整示例可在 MLX GitHub 仓库中找到。