多層感知器

多層感知器#

在此範例中，我們會透過實作簡單的多層感知器來分類 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 class docs。

  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

# 載入資料
import mnist
train_images, train_labels, test_images, test_labels = map(
    mx.array, mnist.mnist()
)

由於我們使用 SGD，因此需要一個迭代器來打亂並建立訓練集的 minibatch：

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 專案中的 完整範例 取得。