|
|
.. _mlp: |
|
|
|
|
|
Multi-Layer Perceptron |
|
|
|
|
|
|
|
|
In this example we'll learn to use ``mlx.nn`` by implementing a simple |
|
|
multi-layer perceptron to classify MNIST. |
|
|
|
|
|
As a first step import the MLX packages we need: |
|
|
|
|
|
.. code-block:: python |
|
|
|
|
|
import mlx.core as mx |
|
|
import mlx.nn as nn |
|
|
import mlx.optimizers as optim |
|
|
|
|
|
import numpy as np |
|
|
|
|
|
|
|
|
The model is defined as the ``MLP`` class which inherits from |
|
|
:class:`mlx.nn.Module`. We follow the standard idiom to make a new module: |
|
|
|
|
|
1. Define an ``__init__`` where the parameters and/or submodules are setup. See |
|
|
the :ref:`Module class docs<module_class>` for more information on how |
|
|
:class:`mlx.nn.Module` registers parameters. |
|
|
2. Define a ``__call__`` where the computation is implemented. |
|
|
|
|
|
.. code-block:: python |
|
|
|
|
|
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) |
|
|
|
|
|
|
|
|
We define the loss function which takes the mean of the per-example cross |
|
|
entropy loss. The ``mlx.nn.losses`` sub-package has implementations of some |
|
|
commonly used loss functions. |
|
|
|
|
|
.. code-block:: python |
|
|
|
|
|
def loss_fn(model, X, y): |
|
|
return mx.mean(nn.losses.cross_entropy(model(X), y)) |
|
|
|
|
|
We also need a function to compute the accuracy of the model on the validation |
|
|
set: |
|
|
|
|
|
.. code-block:: python |
|
|
|
|
|
def eval_fn(model, X, y): |
|
|
return mx.mean(mx.argmax(model(X), axis=1) == y) |
|
|
|
|
|
Next, setup the problem parameters and load the data. To load the data, you need our |
|
|
`mnist data loader |
|
|
<https://github.com/ml-explore/mlx-examples/blob/main/mnist/mnist.py>`_, which |
|
|
we will import as ``mnist``. |
|
|
|
|
|
.. code-block:: python |
|
|
|
|
|
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() |
|
|
) |
|
|
|
|
|
Since we're using SGD, we need an iterator which shuffles and constructs |
|
|
minibatches of examples in the training set: |
|
|
|
|
|
.. code-block:: python |
|
|
|
|
|
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] |
|
|
|
|
|
|
|
|
Finally, we put it all together by instantiating the model, the |
|
|
:class:`mlx.optimizers.SGD` optimizer, and running the training loop: |
|
|
|
|
|
.. code-block:: python |
|
|
|
|
|
# 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}") |
|
|
|
|
|
|
|
|
.. note:: |
|
|
The :func:`mlx.nn.value_and_grad` function is a convenience function to get |
|
|
the gradient of a loss with respect to the trainable parameters of a model. |
|
|
This should not be confused with :func:`mlx.core.value_and_grad`. |
|
|
|
|
|
The model should train to a decent accuracy (about 95%) after just a few passes |
|
|
over the training set. The `full example <https://github.com/ml-explore/mlx-examples/tree/main/mnist>`_ |
|
|
is available in the MLX GitHub repo. |
|
|
|
