web-dev-qa-db-fra.com

Comment éviter "CUDA out of memory" dans PyTorch

Je pense que c'est un message assez courant pour les utilisateurs de PyTorch avec une mémoire GPU faible:

RuntimeError: CUDA out of memory. Tried to allocate ???? MiB (GPU ????; ???? GiB total capacity; ???? GiB already allocated; ???? MiB free; ???? cached)

Je souhaite rechercher des algorithmes de détection d'objets pour mes cours. Et de nombreuses architectures d'apprentissage en profondeur nécessitent une grande capacité de mémoire GPU, de sorte que ma machine ne peut pas entraîner ces modèles. J'ai essayé de traiter une image en chargeant chaque couche sur le GPU, puis en la rechargeant:

for m in self.children():
   m.cuda()
   X = m(X)
   m.cpu()
   torch.cuda.empty_cache()

Mais cela ne semble pas très efficace. Je me demande s'il existe des trucs et astuces pour entraîner de grands modèles d'apprentissage en profondeur tout en utilisant peu de mémoire GPU. Merci d'avance!

Edit: Je suis un débutant en apprentissage profond. Excusez-vous si c'est une question factice :)

2
voilalex

Envoyez les lots à CUDA de manière itérative et créez des lots de petite taille. N'envoyez pas toutes vos données à CUDA en même temps au début. Procédez plutôt comme suit:

for e in range(epochs):
    for images, labels in train_loader:   
        if torch.cuda.is_available():
            images, labels = images.cuda(), labels.cuda()   
        # blablabla  

Vous pouvez également utiliser dtypes qui utilisent moins de mémoire. Par exemple, torch.float16 ou torch.half.

5
Nicolas Gervais

Bien que,

    import torch
    torch.cuda.empty_cache()

fournit une bonne alternative pour effacer la mémoire cuda occupée et nous pouvons également effacer manuellement les variables non utilisées en utilisant,

    import gc
    del variables
    gc.collect()

Mais toujours après avoir utilisé ces commandes, l'erreur peut réapparaître car pytorch n'efface pas réellement la mémoire mais efface la référence à la mémoire occupée par les variables. Donc, réduire le batch_size après avoir redémarré le noyau et trouver le batch_size optimal est la meilleure option possible (mais parfois pas très faisable).

Une autre façon d'obtenir un aperçu plus approfondi de l'alloaction de mémoire dans gpu consiste à utiliser:

    torch.cuda.memory_summary(device=None, abbreviated=False)

où, les deux arguments sont facultatifs. Cela donne un résumé lisible de l'allocation de mémoire et vous permet de comprendre la raison pour laquelle CUDA manque de mémoire et de redémarrer le noyau pour éviter que l'erreur ne se reproduise (comme je l'ai fait dans mon cas).

Passer les données de manière itérative peut aider, mais changer la taille des couches de votre réseau ou les décomposer serait également efficace (car parfois le modèle occupe également une mémoire importante par exemple, lors de l'apprentissage par transfert).

2
SHAGUN SHARMA