La souplesse de yax¶

Les deux premiers tutoriels entraînent des modèles sur des tableaux : une matrice d'entrées X, une matrice de cibles Y, et un objectif de yax.obj. Beaucoup de problèmes ne rentrent pas dans ce moule. yax les accueille sans rien changer à yax.training.train, grâce à trois libertés :

• le modèle est un yax.Module quelconque : sa méthode apply reçoit un exemple et rend ce qu'on veut, un tableau ou un dict ;
• l'objectif est n'importe quelle fonction objective_fn(model, x, y, rkey) qui rend un scalaire ;
• les données X et Y peuvent être des pytrees de tableaux, par exemple des dicts, dont toutes les feuilles ont le même nombre d'exemples en première dimension. yax.training.train les découpe en lots feuille par feuille, et apply reçoit un exemple de même structure.

Cinq cas courants le montrent, chacun sur des données fabriquées : des entrées de natures différentes, des séquences de longueurs différentes, des cibles manquantes, deux sorties de natures différentes, et un objectif qui compare des paires d'exemples.

Télécharger ce notebook. Pour l'ouvrir sur Colab: Fichier → Importer le notebook.

In [ ]:

!pip install -qq -U yaxlib==0.3.*

!pip install -qq -U yaxlib==0.3.*

In [2]:

import shutil
import jax
import jax.numpy as jnp
import jax.random as jr
import optax
import matplotlib.pyplot as plt
import yax

import shutil import jax import jax.numpy as jnp import jax.random as jr import optax import matplotlib.pyplot as plt import yax

In [3]:

dossier = "out/souplesse"
shutil.rmtree(dossier, ignore_errors=True)      # on repart d'un dossier vide
config = yax.configs.TrainConfig(learning_rate=3e-3, nb_epochs=40, lr_final_ratio=0.1)

dossier = "out/souplesse" shutil.rmtree(dossier, ignore_errors=True) # on repart d'un dossier vide config = yax.configs.TrainConfig(learning_rate=3e-3, nb_epochs=40, lr_final_ratio=0.1)

Des entrées de natures différentes¶

Un réseau de vingt magasins vend des glaces. Les ventes d'une journée dépendent du magasin, du jour de la semaine, de la température et d'une éventuelle promotion. Le magasin et le jour sont des catégories : des entiers, qui serviront d'indices. La température et la promotion sont des nombres. Plutôt que de tout convertir en un seul tableau de flottants, on range les entrées dans un dict.

In [5]:

EFFET_MAGASIN = 0.5 * jr.normal(jr.key(0), (20,))          # la loi fabriquée pour l'exemple
EFFET_JOUR = jnp.array([0.0, 0.0, 0.0, 0.0, 0.3, 0.8, 0.6])

def ventes(rkey, nb):
    """nb journées de vente : un dict d'entrées, et les ventes en cible."""
    r = jr.split(rkey, 5)
    magasin = jr.randint(r[0], (nb,), 0, 20)
    jour = jr.randint(r[1], (nb,), 0, 7)
    temperature = jr.uniform(r[2], (nb,), minval=-1, maxval=1)
    promotion = jr.bernoulli(r[3], 0.3, (nb,)).astype(jnp.float32)
    y = EFFET_MAGASIN[magasin] + EFFET_JOUR[jour] + 1.5 * temperature + 0.5 * promotion
    X = {"magasin": magasin, "jour": jour, "nombres": jnp.stack([temperature, promotion], 1)}
    return X, (y + 0.1 * jr.normal(r[4], (nb,)))[:, None]

X1_train, Y1_train = ventes(jr.key(1), 2000)
X1_val, Y1_val = ventes(jr.key(2), 1000)
jax.tree.map(jnp.shape, X1_train)

EFFET_MAGASIN = 0.5 * jr.normal(jr.key(0), (20,)) # la loi fabriquée pour l'exemple EFFET_JOUR = jnp.array([0.0, 0.0, 0.0, 0.0, 0.3, 0.8, 0.6]) def ventes(rkey, nb): """nb journées de vente : un dict d'entrées, et les ventes en cible.""" r = jr.split(rkey, 5) magasin = jr.randint(r[0], (nb,), 0, 20) jour = jr.randint(r[1], (nb,), 0, 7) temperature = jr.uniform(r[2], (nb,), minval=-1, maxval=1) promotion = jr.bernoulli(r[3], 0.3, (nb,)).astype(jnp.float32) y = EFFET_MAGASIN[magasin] + EFFET_JOUR[jour] + 1.5 * temperature + 0.5 * promotion X = {"magasin": magasin, "jour": jour, "nombres": jnp.stack([temperature, promotion], 1)} return X, (y + 0.1 * jr.normal(r[4], (nb,)))[:, None] X1_train, Y1_train = ventes(jr.key(1), 2000) X1_val, Y1_val = ventes(jr.key(2), 1000) jax.tree.map(jnp.shape, X1_train)

Out[5]:

{'jour': (2000,), 'magasin': (2000,), 'nombres': (2000, 2)}

Un lot tiré par yax.training.DatasetSampler a la même structure que X1_train, et ses feuilles gardent leur type : les catégories restent des entiers.

In [7]:

lot_x, lot_y = next(iter(yax.training.DatasetSampler(X1_train, Y1_train, 4)(jr.key(0))))
jax.tree.map(lambda t: f"{t.dtype}{list(t.shape)}", lot_x)

lot_x, lot_y = next(iter(yax.training.DatasetSampler(X1_train, Y1_train, 4)(jr.key(0)))) jax.tree.map(lambda t: f"{t.dtype}{list(t.shape)}", lot_x)

Out[7]:

{'jour': 'int32[4]', 'magasin': 'int32[4]', 'nombres': 'float32[4, 2]'}

Le modèle lit chaque catégorie dans un embedding, une table de vecteurs appris, puis passe ces vecteurs et les nombres à un MLP. Sa méthode apply reçoit un exemple : x["magasin"] est un entier, x["nombres"] un vecteur de deux nombres.

In [9]:

class Ventes(yax.Module):
    magasins: yax.layers.Embedding
    jours: yax.layers.Embedding
    mlp: yax.layers.MLP

    def __init__(self, rkey):
        rkey_m, rkey_j, rkey_mlp = jr.split(rkey, 3)
        self.magasins = yax.layers.Embedding(nb_embeddings=20, dim=4, rkey=rkey_m)
        self.jours = yax.layers.Embedding(nb_embeddings=7, dim=2, rkey=rkey_j)
        self.mlp = yax.layers.MLP(layer_sizes=(4 + 2 + 2, 32, 1), activation="relu", rkey=rkey_mlp)

    def apply(self, x, rkey=None):
        h = jnp.concatenate([self.magasins.apply(x["magasin"]), self.jours.apply(x["jour"]), x["nombres"]])
        return self.mlp.apply(h)

class Ventes(yax.Module): magasins: yax.layers.Embedding jours: yax.layers.Embedding mlp: yax.layers.MLP def __init__(self, rkey): rkey_m, rkey_j, rkey_mlp = jr.split(rkey, 3) self.magasins = yax.layers.Embedding(nb_embeddings=20, dim=4, rkey=rkey_m) self.jours = yax.layers.Embedding(nb_embeddings=7, dim=2, rkey=rkey_j) self.mlp = yax.layers.MLP(layer_sizes=(4 + 2 + 2, 32, 1), activation="relu", rkey=rkey_mlp) def apply(self, x, rkey=None): h = jnp.concatenate([self.magasins.apply(x["magasin"]), self.jours.apply(x["jour"]), x["nombres"]]) return self.mlp.apply(h)

Rien ne change pour l'entraînement : yax.training.train découpe le dict en lots, et yax.obj.mse accepte un x en pytree.

In [11]:

run1 = yax.training.train(f"{dossier}/ventes", config, yax.obj.mse, Ventes(jr.key(0)),
                          (X1_train, Y1_train, 64), (X1_val, Y1_val), rkey=jr.key(3))
print(f"perte de validation {run1.loss:.4f}, pour un bruit de variance 0.01")

run1 = yax.training.train(f"{dossier}/ventes", config, yax.obj.mse, Ventes(jr.key(0)), (X1_train, Y1_train, 64), (X1_val, Y1_val), rkey=jr.key(3)) print(f"perte de validation {run1.loss:.4f}, pour un bruit de variance 0.01")

perte de validation 0.0101, pour un bruit de variance 0.01

⇑ La perte de validation atteint la variance du bruit : le modèle a retrouvé la loi. Les catégories sont restées des entiers jusqu'aux embeddings, sans conversion en flottants.

Des séquences de longueurs différentes¶

Chaque exemple est une liste de 2 à 20 mesures. On veut, pour chaque mesure, son écart à la moyenne de sa liste. Pour ranger des listes de longueurs différentes dans un même tableau, on les complète par des zéros jusqu'à 20 : c'est le remplissage (padding). Un masque dit quelles positions sont de vraies mesures. Il sert deux fois : le modèle ne doit pas compter le remplissage dans sa moyenne, et l'objectif ne doit pas le noter.

In [14]:

T = 20

def listes(rkey, nb):
    """nb listes de 2 à T mesures, complétées par des zéros jusqu'à T."""
    rkey_n, rkey_v = jr.split(rkey)
    longueur = jr.randint(rkey_n, (nb,), 2, T + 1)
    masque = (jnp.arange(T)[None, :] < longueur[:, None]).astype(jnp.float32)
    valeurs = (2 + jr.normal(rkey_v, (nb, T))) * masque
    moyenne = jnp.sum(valeurs, 1, keepdims=True) / jnp.sum(masque, 1, keepdims=True)
    return {"valeurs": valeurs, "masque": masque}, (valeurs - moyenne) * masque

X2_train, Y2_train = listes(jr.key(1), 4000)
X2_val, Y2_val = listes(jr.key(2), 1000)

T = 20 def listes(rkey, nb): """nb listes de 2 à T mesures, complétées par des zéros jusqu'à T.""" rkey_n, rkey_v = jr.split(rkey) longueur = jr.randint(rkey_n, (nb,), 2, T + 1) masque = (jnp.arange(T)[None, :] < longueur[:, None]).astype(jnp.float32) valeurs = (2 + jr.normal(rkey_v, (nb, T))) * masque moyenne = jnp.sum(valeurs, 1, keepdims=True) / jnp.sum(masque, 1, keepdims=True) return {"valeurs": valeurs, "masque": masque}, (valeurs - moyenne) * masque X2_train, Y2_train = listes(jr.key(1), 4000) X2_val, Y2_val = listes(jr.key(2), 1000)

Le modèle transforme chaque mesure en un vecteur, fait la moyenne de ces vecteurs sur les seules vraies mesures, puis rend pour chaque position un nombre calculé à partir de sa mesure et de cette moyenne.

In [16]:

class Ecarts(yax.Module):
    par_mesure: yax.layers.MLP
    sortie: yax.layers.MLP

    def __init__(self, rkey):
        rkey_m, rkey_s = jr.split(rkey)
        self.par_mesure = yax.layers.MLP(layer_sizes=(1, 16, 8), activation="tanh", rkey=rkey_m)
        self.sortie = yax.layers.MLP(layer_sizes=(1 + 8, 16, 1), activation="tanh", rkey=rkey_s)

    def apply(self, x, rkey=None):
        v, m = x["valeurs"][:, None], x["masque"]                      # (T, 1) et (T,)
        h = jax.vmap(self.par_mesure.apply)(v)                         # (T, 8)
        contexte = jnp.sum(h * m[:, None], 0) / jnp.sum(m)             # moyenne sur les vraies mesures
        entree = jnp.concatenate([v, jnp.broadcast_to(contexte, (T, 8))], 1)
        return jax.vmap(self.sortie.apply)(entree)[:, 0]               # (T,) : un écart par position

class Ecarts(yax.Module): par_mesure: yax.layers.MLP sortie: yax.layers.MLP def __init__(self, rkey): rkey_m, rkey_s = jr.split(rkey) self.par_mesure = yax.layers.MLP(layer_sizes=(1, 16, 8), activation="tanh", rkey=rkey_m) self.sortie = yax.layers.MLP(layer_sizes=(1 + 8, 16, 1), activation="tanh", rkey=rkey_s) def apply(self, x, rkey=None): v, m = x["valeurs"][:, None], x["masque"] # (T, 1) et (T,) h = jax.vmap(self.par_mesure.apply)(v) # (T, 8) contexte = jnp.sum(h * m[:, None], 0) / jnp.sum(m) # moyenne sur les vraies mesures entree = jnp.concatenate([v, jnp.broadcast_to(contexte, (T, 8))], 1) return jax.vmap(self.sortie.apply)(entree)[:, 0] # (T,) : un écart par position

L'objectif ne note que les vraies mesures : il lit le masque dans x.

In [18]:

def mse_masquee(model, x, y, rkey=None):
    """L'erreur quadratique moyenne sur les seules vraies mesures."""
    prediction = yax.batch_apply(model, x, rkey)                      # (lot, T)
    return jnp.sum((prediction - y) ** 2 * x["masque"]) / jnp.sum(x["masque"])

run2 = yax.training.train(f"{dossier}/ecarts", config, mse_masquee, Ecarts(jr.key(0)),
                          (X2_train, Y2_train, 64), (X2_val, Y2_val), rkey=jr.key(3))
zero = jnp.sum(Y2_val ** 2 * X2_val["masque"]) / jnp.sum(X2_val["masque"])   # répondre toujours 0
print(f"perte de validation {run2.loss:.4f}, contre {float(zero):.4f} pour un modèle qui répond toujours 0")

def mse_masquee(model, x, y, rkey=None): """L'erreur quadratique moyenne sur les seules vraies mesures.""" prediction = yax.batch_apply(model, x, rkey) # (lot, T) return jnp.sum((prediction - y) ** 2 * x["masque"]) / jnp.sum(x["masque"]) run2 = yax.training.train(f"{dossier}/ecarts", config, mse_masquee, Ecarts(jr.key(0)), (X2_train, Y2_train, 64), (X2_val, Y2_val), rkey=jr.key(3)) zero = jnp.sum(Y2_val ** 2 * X2_val["masque"]) / jnp.sum(X2_val["masque"]) # répondre toujours 0 print(f"perte de validation {run2.loss:.4f}, contre {float(zero):.4f} pour un modèle qui répond toujours 0")

perte de validation 0.0003, contre 0.9064 pour un modèle qui répond toujours 0

Vérifions que le remplissage peut contenir n'importe quoi : remplaçons ses zéros par 99, et comparons les prédictions sur les vraies mesures.

In [20]:

x = jax.tree.map(lambda t: t[0], X2_val)                          # une liste de validation
x_99 = {"valeurs": jnp.where(x["masque"] > 0, x["valeurs"], 99.0), "masque": x["masque"]}
modele2 = run2.trained_model
vraies = x["masque"] > 0
print("longueur de la liste :", int(x["masque"].sum()))
print("écart maximal entre les deux prédictions :", float(jnp.max(jnp.abs(modele2.apply(x) - modele2.apply(x_99))[vraies])))

x = jax.tree.map(lambda t: t[0], X2_val) # une liste de validation x_99 = {"valeurs": jnp.where(x["masque"] > 0, x["valeurs"], 99.0), "masque": x["masque"]} modele2 = run2.trained_model vraies = x["masque"] > 0 print("longueur de la liste :", int(x["masque"].sum())) print("écart maximal entre les deux prédictions :", float(jnp.max(jnp.abs(modele2.apply(x) - modele2.apply(x_99))[vraies])))

longueur de la liste : 5
écart maximal entre les deux prédictions : 0.0

⇑ Les prédictions sur les vraies mesures ne bougent pas : le modèle ne lit le remplissage qu'à travers le masque. Et l'objectif, qui lit le même masque, n'a jamais noté les positions remplies.

Des cibles manquantes¶

Une machine porte trois capteurs, dont on veut prédire les mesures à partir de quatre réglages. Dans les relevés d'entraînement, chaque capteur manque une fois sur deux. Jeter les relevés incomplets ne laisserait qu'un huitième des données. On les garde tous, et la cible devient un dict : les valeurs, et un indicateur de ce qui a vraiment été mesuré.

In [23]:

W = jr.normal(jr.key(0), (4, 3))                     # la loi fabriquée pour l'exemple

def releves(rkey, nb, p_manque):
    """nb relevés ; chaque capteur manque avec la probabilité p_manque."""
    rkey_x, rkey_b, rkey_m = jr.split(rkey, 3)
    X = jr.normal(rkey_x, (nb, 4))
    mesures = jnp.tanh(X @ W) + 0.05 * jr.normal(rkey_b, (nb, 3))
    observee = jr.bernoulli(rkey_m, 1 - p_manque, (nb, 3)).astype(jnp.float32)
    return X, {"valeur": mesures * observee, "observee": observee}

X3_train, Y3_train = releves(jr.key(1), 600, p_manque=0.5)
X3_val, Y3_val = releves(jr.key(2), 1000, p_manque=0.0)      # la validation est complète
print("relevés complets dans l'entraînement :", int(jnp.sum(jnp.all(Y3_train["observee"] > 0, 1))), "sur 600")

W = jr.normal(jr.key(0), (4, 3)) # la loi fabriquée pour l'exemple def releves(rkey, nb, p_manque): """nb relevés ; chaque capteur manque avec la probabilité p_manque.""" rkey_x, rkey_b, rkey_m = jr.split(rkey, 3) X = jr.normal(rkey_x, (nb, 4)) mesures = jnp.tanh(X @ W) + 0.05 * jr.normal(rkey_b, (nb, 3)) observee = jr.bernoulli(rkey_m, 1 - p_manque, (nb, 3)).astype(jnp.float32) return X, {"valeur": mesures * observee, "observee": observee} X3_train, Y3_train = releves(jr.key(1), 600, p_manque=0.5) X3_val, Y3_val = releves(jr.key(2), 1000, p_manque=0.0) # la validation est complète print("relevés complets dans l'entraînement :", int(jnp.sum(jnp.all(Y3_train["observee"] > 0, 1))), "sur 600")

relevés complets dans l'entraînement : 76 sur 600

L'objectif lit les deux feuilles de la cible, et ne compte que les mesures observées. yax.obj.mse ne saurait pas le faire : il attend une cible en tableau.

In [25]:

def mse_observee(model, x, y, rkey=None):
    """L'erreur quadratique moyenne sur les seules mesures observées."""
    prediction = yax.batch_apply(model, x, rkey)                      # (lot, 3)
    return jnp.sum((prediction - y["valeur"]) ** 2 * y["observee"]) / jnp.sum(y["observee"])

def mlp_capteurs():
    return yax.layers.MLP(layer_sizes=(4, 32, 3), activation="tanh", rkey=jr.key(0))

run3 = yax.training.train(f"{dossier}/capteurs", config, mse_observee, mlp_capteurs(),
                          (X3_train, Y3_train, 32), (X3_val, Y3_val), rkey=jr.key(3))

complets = jnp.all(Y3_train["observee"] > 0, 1)               # pour comparer : les seuls relevés complets
run3_complets = yax.training.train(f"{dossier}/capteurs_complets", config, yax.obj.mse, mlp_capteurs(),
                                   (X3_train[complets], Y3_train["valeur"][complets], 32),
                                   (X3_val, Y3_val["valeur"]), rkey=jr.key(3))
print(f"tous les relevés      : perte de validation {run3.loss:.4f}")
print(f"relevés complets seuls : perte de validation {run3_complets.loss:.4f}")

def mse_observee(model, x, y, rkey=None): """L'erreur quadratique moyenne sur les seules mesures observées.""" prediction = yax.batch_apply(model, x, rkey) # (lot, 3) return jnp.sum((prediction - y["valeur"]) ** 2 * y["observee"]) / jnp.sum(y["observee"]) def mlp_capteurs(): return yax.layers.MLP(layer_sizes=(4, 32, 3), activation="tanh", rkey=jr.key(0)) run3 = yax.training.train(f"{dossier}/capteurs", config, mse_observee, mlp_capteurs(), (X3_train, Y3_train, 32), (X3_val, Y3_val), rkey=jr.key(3)) complets = jnp.all(Y3_train["observee"] > 0, 1) # pour comparer : les seuls relevés complets run3_complets = yax.training.train(f"{dossier}/capteurs_complets", config, yax.obj.mse, mlp_capteurs(), (X3_train[complets], Y3_train["valeur"][complets], 32), (X3_val, Y3_val["valeur"]), rkey=jr.key(3)) print(f"tous les relevés : perte de validation {run3.loss:.4f}") print(f"relevés complets seuls : perte de validation {run3_complets.loss:.4f}")

tous les relevés      : perte de validation 0.0514
relevés complets seuls : perte de validation 0.1216

⇑ Avec tous les relevés, la perte de validation est plus de deux fois plus basse qu'avec les seuls relevés complets. Les mesures partielles apportent de l'information, à condition que l'objectif sache les lire.

Deux sorties de natures différentes¶

Un point du plan appartient à l'un de trois secteurs angulaires, et l'on veut aussi sa distance à l'origine : une classe et un nombre. Un seul modèle, deux têtes : apply rend un dict. La cible est un dict de même structure, et l'objectif additionne l'entropie croisée de la classe et l'erreur quadratique de la distance.

In [28]:

def points(rkey, nb):
    """nb points du plan, leur secteur (0, 1 ou 2) et leur distance à l'origine."""
    rkey_r, rkey_a = jr.split(rkey)
    rayon = jr.uniform(rkey_r, (nb,), minval=0.2, maxval=2.0)
    angle = jr.uniform(rkey_a, (nb,), maxval=2 * jnp.pi)
    X = jnp.stack([rayon * jnp.cos(angle), rayon * jnp.sin(angle)], 1)
    secteur = jnp.floor(angle / (2 * jnp.pi / 3)).astype(jnp.int32)
    return X, {"classe": secteur, "distance": rayon[:, None]}

X4_train, Y4_train = points(jr.key(1), 3000)
X4_val, Y4_val = points(jr.key(2), 1000)

class DeuxTetes(yax.Module):
    corps: yax.layers.MLP
    tete_classe: yax.layers.Linear
    tete_distance: yax.layers.Linear

    def __init__(self, rkey):
        rkey_c, rkey_tc, rkey_td = jr.split(rkey, 3)
        self.corps = yax.layers.MLP(layer_sizes=(2, 64, 64), activation="relu", rkey=rkey_c)
        self.tete_classe = yax.layers.Linear(dim_in=64, dim_out=3, rkey=rkey_tc)
        self.tete_distance = yax.layers.Linear(dim_in=64, dim_out=1, rkey=rkey_td)

    def apply(self, x, rkey=None):
        h = jax.nn.relu(self.corps.apply(x))              # le MLP n'active pas sa sortie
        return {"classe": self.tete_classe.apply(h), "distance": self.tete_distance.apply(h)}

def points(rkey, nb): """nb points du plan, leur secteur (0, 1 ou 2) et leur distance à l'origine.""" rkey_r, rkey_a = jr.split(rkey) rayon = jr.uniform(rkey_r, (nb,), minval=0.2, maxval=2.0) angle = jr.uniform(rkey_a, (nb,), maxval=2 * jnp.pi) X = jnp.stack([rayon * jnp.cos(angle), rayon * jnp.sin(angle)], 1) secteur = jnp.floor(angle / (2 * jnp.pi / 3)).astype(jnp.int32) return X, {"classe": secteur, "distance": rayon[:, None]} X4_train, Y4_train = points(jr.key(1), 3000) X4_val, Y4_val = points(jr.key(2), 1000) class DeuxTetes(yax.Module): corps: yax.layers.MLP tete_classe: yax.layers.Linear tete_distance: yax.layers.Linear def __init__(self, rkey): rkey_c, rkey_tc, rkey_td = jr.split(rkey, 3) self.corps = yax.layers.MLP(layer_sizes=(2, 64, 64), activation="relu", rkey=rkey_c) self.tete_classe = yax.layers.Linear(dim_in=64, dim_out=3, rkey=rkey_tc) self.tete_distance = yax.layers.Linear(dim_in=64, dim_out=1, rkey=rkey_td) def apply(self, x, rkey=None): h = jax.nn.relu(self.corps.apply(x)) # le MLP n'active pas sa sortie return {"classe": self.tete_classe.apply(h), "distance": self.tete_distance.apply(h)}

yax.batch_apply applique le modèle à chaque exemple du lot, et rend un dict de lots. L'objectif a une option, le poids de la distance : le décorateur yax.obj.objective la rend réglable à l'appel, comme les objectifs de yax.obj.

In [30]:

@yax.obj.objective
def classe_et_distance(model, x, y, rkey=None, *, poids=1.0):
    """Entropie croisée de la classe, plus poids fois l'erreur quadratique de la distance."""
    sortie = yax.batch_apply(model, x, rkey)                         # {"classe": (lot, 3), "distance": (lot, 1)}
    ce = jnp.mean(optax.softmax_cross_entropy_with_integer_labels(sortie["classe"], y["classe"]))
    mse = jnp.mean((sortie["distance"] - y["distance"]) ** 2)
    return ce + poids * mse

run4 = yax.training.train(f"{dossier}/deux_tetes", config, classe_et_distance(poids=1.0), DeuxTetes(jr.key(0)),
                          (X4_train, Y4_train, 64), (X4_val, Y4_val), rkey=jr.key(3))
sortie = yax.batch_apply(run4.trained_model, X4_val)
justesse = jnp.mean(jnp.argmax(sortie["classe"], -1) == Y4_val["classe"])
rmse = jnp.sqrt(jnp.mean((sortie["distance"] - Y4_val["distance"]) ** 2))
print(f"secteur : justesse {float(justesse):.3f}   distance : RMSE {float(rmse):.3f}")

@yax.obj.objective def classe_et_distance(model, x, y, rkey=None, *, poids=1.0): """Entropie croisée de la classe, plus poids fois l'erreur quadratique de la distance.""" sortie = yax.batch_apply(model, x, rkey) # {"classe": (lot, 3), "distance": (lot, 1)} ce = jnp.mean(optax.softmax_cross_entropy_with_integer_labels(sortie["classe"], y["classe"])) mse = jnp.mean((sortie["distance"] - y["distance"]) ** 2) return ce + poids * mse run4 = yax.training.train(f"{dossier}/deux_tetes", config, classe_et_distance(poids=1.0), DeuxTetes(jr.key(0)), (X4_train, Y4_train, 64), (X4_val, Y4_val), rkey=jr.key(3)) sortie = yax.batch_apply(run4.trained_model, X4_val) justesse = jnp.mean(jnp.argmax(sortie["classe"], -1) == Y4_val["classe"]) rmse = jnp.sqrt(jnp.mean((sortie["distance"] - Y4_val["distance"]) ** 2)) print(f"secteur : justesse {float(justesse):.3f} distance : RMSE {float(rmse):.3f}")

secteur : justesse 0.998   distance : RMSE 0.010

⇑ Les deux têtes apprennent ensemble, sur un seul objectif, et réussissent toutes les deux. L'option poids règle l'équilibre entre les deux tâches : classe_et_distance(poids=10.0) donnerait dix fois plus d'importance à la distance.

Des paires d'exemples¶

Des points de dimension 6 appartiennent à quatre classes, mais on ne connaît pas ces classes. On sait seulement, pour des paires de points, s'ils sont de la même classe. On veut une représentation en deux dimensions où les points d'une même classe sont proches.

Le modèle s'applique à un point, et rend sa représentation. C'est l'objectif qui l'applique aux deux points de chaque paire, et compare leurs représentations : plus elles sont proches, plus la paire est jugée de même classe. Le seuil de cette décision est un paramètre du modèle, appris avec le reste.

In [33]:

CENTRES = 2.0 * jr.normal(jr.key(0), (4, 6))                   # quatre classes cachées

def nuage(rkey, nb):
    rkey_c, rkey_b = jr.split(rkey)
    classe = jr.randint(rkey_c, (nb,), 0, 4)
    return CENTRES[classe] + jr.normal(rkey_b, (nb, 6)), classe

def paires(rkey, nb):
    """nb paires de points, dont une moitié environ de même classe."""
    rkey_a, rkey_b, rkey_c, rkey_m = jr.split(rkey, 4)
    a, classe_a = nuage(rkey_a, nb)
    b, classe_b = nuage(rkey_b, nb)
    meme = jr.bernoulli(rkey_m, 0.5, (nb,))                       # forcer la même classe une fois sur deux
    b = jnp.where(meme[:, None], CENTRES[classe_a] + jr.normal(rkey_c, (nb, 6)), b)
    meme = meme | (classe_a == classe_b)
    return {"a": a, "b": b}, meme.astype(jnp.float32)[:, None]

X5_train, Y5_train = paires(jr.key(1), 4000)
X5_val, Y5_val = paires(jr.key(2), 1000)

class Plongement(yax.Module):
    mlp: yax.layers.MLP
    seuil: jax.Array                                                # un paramètre scalaire

    def __init__(self, rkey):
        self.mlp = yax.layers.MLP(layer_sizes=(6, 32, 2), activation="relu", rkey=rkey)
        self.seuil = jnp.array(1.0)

    def apply(self, x, rkey=None):
        return self.mlp.apply(x)                                    # un point -> sa représentation

def meme_classe(model, x, y, rkey=None):
    """Deux points proches dans la représentation sont jugés de même classe."""
    za = yax.batch_apply(model, x["a"])
    zb = yax.batch_apply(model, x["b"])
    logit = model.seuil - jnp.sum((za - zb) ** 2, -1, keepdims=True)
    return jnp.mean(optax.sigmoid_binary_cross_entropy(logit, y))

run5 = yax.training.train(f"{dossier}/paires", config, meme_classe, Plongement(jr.key(0)),
                          (X5_train, Y5_train, 64), (X5_val, Y5_val), rkey=jr.key(3))
print(f"perte de validation {run5.loss:.4f}")

CENTRES = 2.0 * jr.normal(jr.key(0), (4, 6)) # quatre classes cachées def nuage(rkey, nb): rkey_c, rkey_b = jr.split(rkey) classe = jr.randint(rkey_c, (nb,), 0, 4) return CENTRES[classe] + jr.normal(rkey_b, (nb, 6)), classe def paires(rkey, nb): """nb paires de points, dont une moitié environ de même classe.""" rkey_a, rkey_b, rkey_c, rkey_m = jr.split(rkey, 4) a, classe_a = nuage(rkey_a, nb) b, classe_b = nuage(rkey_b, nb) meme = jr.bernoulli(rkey_m, 0.5, (nb,)) # forcer la même classe une fois sur deux b = jnp.where(meme[:, None], CENTRES[classe_a] + jr.normal(rkey_c, (nb, 6)), b) meme = meme | (classe_a == classe_b) return {"a": a, "b": b}, meme.astype(jnp.float32)[:, None] X5_train, Y5_train = paires(jr.key(1), 4000) X5_val, Y5_val = paires(jr.key(2), 1000) class Plongement(yax.Module): mlp: yax.layers.MLP seuil: jax.Array # un paramètre scalaire def __init__(self, rkey): self.mlp = yax.layers.MLP(layer_sizes=(6, 32, 2), activation="relu", rkey=rkey) self.seuil = jnp.array(1.0) def apply(self, x, rkey=None): return self.mlp.apply(x) # un point -> sa représentation def meme_classe(model, x, y, rkey=None): """Deux points proches dans la représentation sont jugés de même classe.""" za = yax.batch_apply(model, x["a"]) zb = yax.batch_apply(model, x["b"]) logit = model.seuil - jnp.sum((za - zb) ** 2, -1, keepdims=True) return jnp.mean(optax.sigmoid_binary_cross_entropy(logit, y)) run5 = yax.training.train(f"{dossier}/paires", config, meme_classe, Plongement(jr.key(0)), (X5_train, Y5_train, 64), (X5_val, Y5_val), rkey=jr.key(3)) print(f"perte de validation {run5.loss:.4f}")

perte de validation 0.0241

Les classes n'ont jamais été données au modèle. Dessinons pourtant la représentation de points neufs, colorés selon leur classe :

In [35]:

points_neufs, classes = nuage(jr.key(9), 800)
z = yax.batch_apply(run5.trained_model, points_neufs)
_, ax = plt.subplots(figsize=(5, 4))
ax.scatter(z[:, 0], z[:, 1], c=classes, cmap="tab10", vmin=0, vmax=9, s=8)
ax.set_title("la représentation apprise, colorée par la classe cachée", fontsize=10)
plt.show()

points_neufs, classes = nuage(jr.key(9), 800) z = yax.batch_apply(run5.trained_model, points_neufs) _, ax = plt.subplots(figsize=(5, 4)) ax.scatter(z[:, 0], z[:, 1], c=classes, cmap="tab10", vmin=0, vmax=9, s=8) ax.set_title("la représentation apprise, colorée par la classe cachée", fontsize=10) plt.show()

⇑ Quatre groupes nets, un par classe : la représentation a retrouvé les classes à partir des seules paires. Les couleurs viennent des classes cachées, que l'entraînement n'a jamais vues.

Et ensuite ?¶

• yax.obj : les objectifs fournis, et le décorateur yax.obj.objective qui rend les vôtres réglables.
• yax.training : yax.training.train, DatasetSampler et FunctionSampler, et ce qu'ils acceptent comme données.
• yax.layers : les briques, dont Embedding et les couches récurrentes, qui lisent des séquences.