GPU

Si certains d’entre vous tentent l’expérience GPU, vous allez probablement rencontré le message d’erreur : out of memory (OOM), même sur de petits datasets.

Ce comportement peut paraître étrange. Nous allons donc expliquer d’où vient ce problème et comment l’interpréter.

Consommation mémoire

Il y a trois sources de consommation mémoire

Poids

En exécutant le code suivant, nous pouvons déterminer rapidement combien de poids sont présents dans le réseau :

sum(p.numel() for p in model.parameters())
>> 538 981

Pour chaque poids dans le réseau, nous devons stocker :

  • Sa valeur

  • La valeur de son gradient

  • Deux valeurs supplémentaires pour l’optimiseur Adam (et oui !)

Un float32 faisant 4 octets, avec 4 valeurs par poids, cela fait un total de :

  • Nombre de poids x 16 octects

Ainsi, pour notre réseau avec 500k poids, cela fait : 8Mo.

Batch

Chaque image représente 3x128x128x4 octets soit 200ko. Il faut ensuite tenir compte de la taille du batch :

  • 32 images : 6Mo

  • 256 images : 50Mo

Activation

Pour calculer la rétropropagation, il faut conserver les sorties de chaque couche. Si l’on regarde uniquement les couches Conv2D, nous avons :

  • conv2D_1 : 32 x 128 x 128

  • conv2D_2 : 64 x 64 x 64

  • conv2D_3 : 128 x 32 x 32

  • conv2D_4 : 256

Soit un total de ~917 000 valeurs, donc en float32 : 3,6 Mo/image. Ainsi, pour stocker les activations :

  • avec un batch de 32 : 115 Mo

  • avec un batch de 256 : 921 Mo

Cette grandeur est sous-estimée, car il faut aussi comptabiliser les buffers annexes qui rajoutent 30% à 60%.

Conclusion

Dans notre modèle avec un batch de 32 images, nous obtenons un total de :

  • Poids : 8 Mo

  • Batch : 6 Mo

  • Activation + Buffer : ~160Mo

  • Total : 174 Mo

Ainsi, c’est le stockage des activations qui coûtent le plus cher et c’est lui qui provoque un pic de consommation mémoire.

Attention donc à la taille des batchs.

Dans notre cas, multiplier la taille du batch par 8 (256), multiplie d’autant la mémoire totale :

  • Poids : 8 Mo

  • Batch : 50 Mo

  • Activation + Buffer : ~1400 Mo

  • Total : 1458 Mo

A celà, il faut ajouter des buffers de calcul qui ajoutent 20% à 40% de taille mémoire supplémentaire.

Note

Lorsque vous gérez un modèle professionnel, la question de la taille mémoire devient cruciale. Autant pour une inférence, la taille peut paraître raisonable, autant durant l’apprentissage, il faut faire extrêmement attention.