Les librairies de deep learning sont très conservatrices dans le choix des types de données.
Par défaut, elles utilisent des types 64-bit ou 32-bit, alors qu'un 16-bit serait souvent suffisant. Tu peux d'ailleurs le constater dans le code ci-dessous :
Résultat : on est loin d'une allocation mémoire optimale.
C'est un choix délibéré, fait pour garantir une meilleure précision dans la représentation des valeurs numériques.
Mais cette précision a un coût : une consommation mémoire plus élevée, ce qui n'est pas toujours souhaitable.
On observe aussi que beaucoup d'opérations tensorielles — notamment les multiplications matricielles — sont bien plus rapides avec des types de faible précision, comme le montre l'image suivante :
De plus, comme le float16 ne pèse que la moitié d'un float32, l'utiliser réduit la mémoire nécessaire à l'entraînement.
Ça permet aussi d'entraîner des modèles plus grands, d'utiliser des mini-batches plus importants (et donc d'accélérer encore davantage), etc.
Le mixed precision training est une technique fiable et largement adoptée dans l'industrie pour y parvenir.
Comme son nom l'indique, l'idée est de combiner une précision réduite en float16 (là où c'est faisable, comme dans les convolutions et les multiplications matricielles) avec du float32 — d'où le terme "mixed precision".
Voici une liste de modèles entraînés avec le mixed precision training :
C'est clairement une pratique très répandue… mais dont on parle peu.
Avant d'entrer dans les détails techniques…
Ce qui précède doit te donner une intuition : en utilisant un type de faible précision (float16), on peut introduire des incohérences et des imprécisions numériques sans le vouloir.
Pour les éviter, voici les bonnes pratiques du mixed precision training — avec le code associé.
Mixed precision training avec PyTorch : bonnes pratiques
Utiliser le mixed precision training avec PyTorch nécessite quelques ajustements dans ton implémentation d'entraînement.
Voici une implémentation classique d'entraînement PyTorch :
La première chose qu'on introduit, c'est un objet scaler qui va mettre à l'échelle la valeur de la loss :
Pourquoi ?
Parce que la valeur originale de la loss peut parfois être si faible qu'il devient impossible de calculer les gradients en float16 avec une précision suffisante.
Dans ce cas, aucune mise à jour des poids du modèle ne se produit.
Mettre la loss à l'échelle vers une plage numérique plus élevée garantit que même les petits gradients contribuent à la mise à jour des poids.
Mais ces petits gradients ne peuvent être intégrés dans la matrice de poids que si celle-ci est représentée en haute précision — c'est-à-dire en float32.
C'est pourquoi, par mesure de précaution, on conserve les poids en float32.
Cela dit, le scaling de la loss n'est pas toujours indispensable : dans mon expérience, ces petites mises à jour n'apparaissent généralement qu'en fin d'entraînement.
On peut donc raisonnablement supposer qu'elles n'impactent pas vraiment les performances du modèle.
Mais ne prends pas ça pour une conclusion définitive — c'est quelque chose à valider toi-même quand tu utiliseras le mixed precision training.
Revenons à nos moutons.
Puisque les poids (qui sont des matrices) sont représentés en float32, on ne peut pas bénéficier des accélérations liées au float16 tant qu'ils restent dans cet état :
Voici les étapes à suivre pour exploiter ces accélérations :
On crée une copie en
float16des poids pendant la passe forward.On calcule la valeur de la loss en
float32, puis on la met à l'échelle pour avoir plus de précision dans les gradients, qui sont calculés enfloat16.Les gradients sont calculés en
float16car, comme la passe forward, le calcul des gradients implique aussi des multiplications matricielles.Les conserver en
float16apporte donc un gain de vitesse supplémentaire.
Une fois les gradients calculés en
float16, les opérations matricielles lourdes sont terminées. Il ne reste plus qu'à mettre à jour la matrice de poids originale, qui est enfloat32.On crée donc une copie en
float32des gradients ci-dessus, on retire le scaling appliqué à l'étape 2, et on met à jour les poids enfloat32.Terminé !
Les paramètres de mixed precision pour la passe forward sont gérés par le context manager torch.autocast() :
Place maintenant à la passe backward.
Ligne 13 →
scaler.scale(loss).backward(): l'objetscalermet la loss à l'échelle, puisbackward()est appelé pour calculer les gradients.Ligne 14 →
scaler.step(opt): dé-mise à l'échelle des gradients et mise à jour des poids.Ligne 15 →
scaler.update(): mise à jour du scale pour l'itération suivante.Ligne 16 →
opt.zero_grad(): remise à zéro des gradients.
Et voilà !
L'efficacité du mixed precision scaling par rapport à un entraînement classique est bien visible sur l'image suivante :
Le mixed precision training est plus de 2,5× plus rapide qu'un entraînement conventionnel.
Plutôt sympa, non ?
Pour plus de détails sur le code, consulte la documentation officielle PyTorch : Documentation PyTorch — Mixed Precision.
Une autre façon efficace d'accélérer l'entraînement d'un modèle : le Momentum. On en a parlé récemment dans ce shot data : Le Momentum en ML, expliqué visuellement.
Envie d'aller plus loin ? DataXp t'attend : formations et parcours pour progresser en data science, machine learning et IA.
👉 À toi de jouer : quelles autres techniques fiables connais-tu pour accélérer l'entraînement de tes modèles de machine learning ?
