Chargement de ResNet50 sur RTX2070 - Manque de mémoire

Question

J'essaie de charger ResNext50, et en plus CenterNet, je suis capable de le faire avec le GPU de Google Colab ou de Kaggle. Mais,

1. J'aimerais savoir de combien de mémoire GPU (VRAM) ce réseau a-t-il besoin ?

2. Lorsque j'utilise une RTX 2070 avec 5,5 Go de VRAM libre (sur 8 Go), je n'arrive pas à le charger.

La taille du lot est de 1, le nombre de travailleurs est de 1, tout est réglé sur les valeurs minimales. SYSTÈME D'EXPLOITATION : Ubuntu 18.04 (Utilisation de PyTorch)

Dans TensorFlow, je sais que je peux restreindre la quantité de VRAM (ce qui me permet de charger et d'exécuter des réseaux bien que je n'aie pas assez de VRAM), mais dans PyTorch je n'ai pas encore trouvé cette fonctionnalité.

Une idée pour résoudre ce problème ?

Answer 1

Utilisation de la dépendance d'un tiers

Vous pourriez obtenir la taille de model en octets en utilisant une bibliothèque tierce torchfunc (avis de non-responsabilité, je suis l'auteur).

import torchfunc

# Assuming model is loaded
print(torchfunc.sizeof(model))

Aucune dépendance

Cette fonction est assez simple et chemise, vous pourriez simplement la copier, voir le code source

Taille exacte

C'est juste la taille de votre modèle, il y a une plus grande utilisation de la mémoire VRAM pendant forward y backward et dépend de la taille de votre batch . Vous pouvez essayer pytorch_modelsize pour les estimer (je ne suis pas sûr que cela fonctionne pour vos réseaux).

Effacer le cache

Vous devriez vider votre cache avant d'exécuter le réseau (parfois le redémarrage de la station de travail a aidé dans mon cas), car vous avez certainement assez de mémoire. Tuer les processus en utilisant le GPU devrait également aider.

Encore une fois torchfunc pourrait vous aider, lancez la commande suivante :

import torchfunc

torchfunc.cuda.reset()