MNIST avec MLX 🤗 (en français)

Cette manip est un peu commme le hello word de l’IA.

Pour moi c’est aussi la première fois que je touche a MLX : le framework IA d’Apple.

Je connais déjà Tensorflow/Keras de Google qui est en train d’être remplacé par JAX/Flax.

J’utilise MLX car c’est un framework élégant et performant, et surtout parce que j’ai un Mac et que c’est la solution la plus adaptée pour exploiter au mieux mon materiel.

Dataset

On télécharge MNIST via la bibliotèque datasets d’HuggingFace 🤗 (merci à eux de faire de l’open-source)/

On utilise le format numpy car datasets ne peut pas (encore ?) fournir directement des tenseurs au format MLX.

from datasets import load_dataset

import mlx.core as mx

dataset = load_dataset("ylecun/mnist").with_format("numpy")

# On split
dataset_train = dataset['train']
dataset_test = dataset['test']

# Pour les données de tests, on les convertie directement 
# comme on veut  et on instancie le tensur MLX
test_images = mx.array(dataset_test['image'].reshape((-1, 28*28)))
test_labels = mx.array(dataset_test['label'])

test_images.shape, test_labels.shape

((10000, 784), (10000,))

from IPython.display import display
import PIL
display(PIL.Image.fromarray(dataset_train[0]['image']))

Donc le dataset de test contient 10000 pair image-label.

Les images ont des dimention de 784 (28x28) et les labels sont des scalairs (pas de dimention). J’ai applati les images car le model sera un Perceptron_multicouche (voir plus bas).

Pour l’apprentissage, on veut créer des catégories pour chaque chiffre à deviner. Il faut donc faire le un encodage OneHot.

import numpy as np

def onehot10(labels):
    b = np.zeros((labels.size, 10))
    b[np.arange(labels.size), labels] = 1
    return b

# Cette fonction est un midleware pour la lib datasets. s'applique sur des batch
def transforms_onehot10(data):
    data['image'] = mx.array(np.array(data['image']).reshape(-1, 28*28) / 255)
    data['label'] = mx.array(onehot10(np.array(data['label'])))
    return data

dataset_train.set_transform(transforms_onehot10)

Model

On on fait un MLP ou Perceptron_multicouche.

Voici pour une couche de perceptron : $y = relu(W x + B)$

avec

• $W$ (weight) une matrice pas forcément carré
• $B$ (bias) un vecteur de la taille de la sortie $y$
• $relu()$ notre fonction d’activation. C’est surtout pour rendre notre nodel non lineaire.

Notre derniere couche n’a pas de fonction d’activation intégré au model car il n’y en a pas besoin :

• Pour l’apprentissage, la fonction de coup utilisée sera cross_entropy() et il ne faut pas normaliser les données
• Pour l’évaluation, on regarde simplement la classe qui a la plus grande valeure

from functools import reduce

import mlx.nn as nn


class MNIST_MLP(nn.Module):
    def __init__(self):
        super().__init__()
        self.layers = [
            nn.Linear(28 * 28, 32), # Aplique Wx+B
            nn.relu,
            nn.Linear(32, 32),
            nn.relu,
            nn.Linear(32, 10),
        ]

    def __call__(self, x):
        # En gros on compose les fonctions de self.layers
        return reduce(lambda v, f: f(v), self.layers, x)

model = MNIST_MLP()
model

MNIST_MLP(
  (layers.0): Linear(input_dims=784, output_dims=32, bias=True)
  (layers.2): Linear(input_dims=32, output_dims=32, bias=True)
  (layers.4): Linear(input_dims=32, output_dims=10, bias=True)
)

Fonction de cout, optimiseur et évaluateur

On instancie d’abore un SGD (Stochastic_gradient_descent) qui est l’algorithme qui permet de trouver les paramètres minimisant notre fonction de coup.

Ensuite, je définie une la fonction de coup : elle applique le model model sur les entrées X puis on calcul l’écart avec les données attendues y avec la méthode cross_entropy. Cette fonction est dérivable et c’est celle ci que l’optimiseur cherche a minimiser.

La fonction d’évaluation permet de calculer la ratio où le model s’est trompé. Ce n’est pas la meme chose que la fonction de coup car celle ci n’est pas dérivable et elle sert juste pour avoir une idée de la qualité du model. On l’applique sur les données du jeu de test.

Si l’évaluation sur les données de testes decroit lors de l’apprentissage, c’est probablement du a un sur-apprentissage.

Vu que X et y sont des batch de données, on moyenne le tout a chaque évalutation.

from mlx.optimizers import SGD

optimizer = SGD(learning_rate=0.1)

def loss_fn(model, X, y):
    return nn.losses.cross_entropy(model(X), y, reduction="mean")

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

Étape d’apprentissage

A chaque étape d’apprentissage, on choisie un batch de donnée du dataset d’apprentissage. On applique la fonction de coup loss_fn() et on calcule sont gadient par rapport au poids du model que l’on veut optimiser.

Ainsi, pour les données du batch en cours, on sait l’effet de chaque poid du model sur la valeur de la fonciton de coup. On sait donc quels poids on beaucoup d’impact (positif et négatif) et on est en mesure d’améliorer le model.

MLX propose une foncion qui fait tout ca d’un coup : nn.value_and_grad()

# loss_fn prend comme arguments (model, X, y)
# loss_and_grad_fn prendra donc aussi les arguments (model, X, y)
# Elle renvera la valeure de loss_fn 
#   ainsi que les gradient par rapport aux poids entrainables.
loss_and_grad_fn = nn.value_and_grad(model, loss_fn)

def step(X, y):
    loss, grads = loss_and_grad_fn(model, X, y)
    # On optimise le model avec les gradient obtenus.
    optimizer.update(model, grads)
    # On renvoie la loss pour faire un suivit de l'apprentissage si on veut.
    return loss

Apprentissage

Pour l’apprentissage, on va passer tout le dataset plusieurs fois car on n’a as une infinité de donnée. C’est comme s’entrainer a conduire sur un circuit. On fait le meme circuit plusieurs fois 🚗. On appelle ca epoch

import time

for epoch in range(5):
    tic = time.perf_counter() # Pour chronometrer le temps d'une epoch.

    # On melange le dataset d'apprentissage et on iter dessus 
    #   en faisant des batchs de 128 couple photo-label.
    for batch in dataset_train.shuffle().iter(128):
        step(batch['image'], batch['label'])  # mx.eval(model.state)

    # A la fin de l'époch, on évalue le model.
    accuracy = eval_fn(test_images, test_labels)
    print(f"Epoch {epoch}: Test accuracy {accuracy.item():.3f}, "
          f"Time {time.perf_counter() - tic:.3f} (s)")

Epoch 0: Test accuracy 0.893, Time 1.965 (s)
Epoch 1: Test accuracy 0.917, Time 1.451 (s)
Epoch 2: Test accuracy 0.923, Time 1.421 (s)
Epoch 3: Test accuracy 0.934, Time 1.356 (s)
Epoch 4: Test accuracy 0.933, Time 1.371 (s)