Aller au contenu

yax.optimizers

Les optimiseurs disponibles pour l'entraînement.

Optimiseurs, désignés par leur nom dans yax.configs.TrainConfig.

Chaque nom de OPTIMIZERS correspond à un constructeur d'optax ("adam" correspond à optax.adam). yax.training.train l'appelle avec le pas d'apprentissage et les optimizer_options de la configuration. On peut aussi passer le constructeur lui-même : optimizer=yax.optimizers.adam équivaut à optimizer="adam".

Optimiseurs disponibles : adam, adamw, adamax, adabelief, adagrad, nadam, radam, rmsprop, lion, sgd et lbfgs.

OPTIMIZERS module-attribute

yax.optimizers.OPTIMIZERS = {'adam': adam, 'adamw': adamw, 'adamax': adamax, 'adabelief': adabelief, 'adagrad': adagrad, 'nadam': nadam, 'radam': radam, 'rmsprop': rmsprop, 'lion': lion, 'sgd': sgd, 'lbfgs': lbfgs}

resolve_optimizer

yax.optimizers.resolve_optimizer(optimizer)

Renvoie le constructeur d'optimiseur désigné par optimizer.

Paramètres :

Nom Type Description Défaut
optimizer str | Callable

un nom de OPTIMIZERS, ou un constructeur optax, renvoyé tel quel.

obligatoire

Lève :

Type Description
ValueError

si le nom est inconnu.

Code source dans yax/optimizers.py
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
def resolve_optimizer(optimizer: str | Callable) -> Callable:
    """Renvoie le constructeur d'optimiseur désigné par `optimizer`.

    Args:
        optimizer: un nom de `OPTIMIZERS`, ou un constructeur optax, renvoyé tel quel.

    Raises:
        ValueError: si le nom est inconnu.
    """
    if callable(optimizer):
        return optimizer
    if optimizer in OPTIMIZERS:
        return OPTIMIZERS[optimizer]
    raise ValueError(
        f"optimiseur inconnu : {optimizer!r}. Donnez un constructeur optax, ou "
        f"une chaîne parmi : {', '.join(OPTIMIZERS)}.")

lbfgs

yax.optimizers.lbfgs(learning_rate=None, **options)

L-BFGS : méthode quasi-Newton avec recherche linéaire.

Le pas est choisi à chaque itération par la recherche linéaire : le learning_rate de la configuration est ignoré. Deux conditions d'emploi :

  • la perte doit être la même à chaque pas : un seul lot fixe, qui couvre tout le problème (données passées en (X, Y), sans batch_size) ;
  • la mémoire utilisée vaut environ vingt copies du modèle.

Adapté pour affiner un modèle déjà entraîné, ou pour les problèmes sans bruit (méthode de Ritz, PINN).

Paramètres :

Nom Type Description Défaut
learning_rate Any

ignoré.

None
**options Any

options de optax.lbfgs, par exemple memory_size.

{}
Code source dans yax/optimizers.py
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
def lbfgs(learning_rate: Any = None, **options: Any) -> optax.GradientTransformationExtraArgs:
    """L-BFGS : méthode quasi-Newton avec recherche linéaire.

    Le pas est choisi à chaque itération par la recherche linéaire : le
    `learning_rate` de la configuration est ignoré. Deux conditions d'emploi :

    - la perte doit être la même à chaque pas : un seul lot fixe, qui couvre
      tout le problème (données passées en `(X, Y)`, sans `batch_size`) ;
    - la mémoire utilisée vaut environ vingt copies du modèle.

    Adapté pour affiner un modèle déjà entraîné, ou pour les problèmes sans
    bruit (méthode de Ritz, PINN).

    Args:
        learning_rate: ignoré.
        **options: options de `optax.lbfgs`, par exemple `memory_size`.
    """
    return optax.lbfgs(learning_rate=None, **options)