Un laboratoire de gestion de la mémoire CUDA simple et précis pour Pytorch, il se compose de différentes parties de la mémoire:
Caractéristiques:
line_profiler style CUDA Memory Profiler avec API simple.%mlrun / %%mlrun LINE / CELL Magic.Table des matières
pip install pytorch_memlabpip install git+https://github.com/stonesjtu/pytorch_memlabLes erreurs hors mémoire dans le pytorch se produisent fréquemment, pour les nouveaux braises et les programmeurs expérimentés. Une raison courante est que la plupart des gens n'apprennent pas vraiment la philosophie de gestion de la mémoire sous-jacente de Pytorch et GPU. Ils ont écrit des codes de mémoire et se sont plaints que Pytorch mange trop de mémoire Cuda.
Dans ce dépôt, je vais partager des outils utiles pour aider à déboguer OOM, ou pour inspecter le mécanisme sous-jacent si quelqu'un s'intéresse.
Le profileur de mémoire est une modification de line_profiler de Python, il donne les informations d'utilisation de la mémoire pour chaque ligne de code dans la fonction / méthode spécifiée.
import torch
from pytorch_memlab import LineProfiler
def inner ():
torch . nn . Linear ( 100 , 100 ). cuda ()
def outer ():
linear = torch . nn . Linear ( 100 , 100 ). cuda ()
linear2 = torch . nn . Linear ( 100 , 100 ). cuda ()
linear3 = torch . nn . Linear ( 100 , 100 ). cuda ()
work ()Une fois le script terminé ou interrompu par le clavier, il donne les informations de profilage suivantes si vous êtes dans un cahier Jupyter:

ou les informations suivantes si vous êtes dans un terminal en texte uniquement:
## outer
active_bytes reserved_bytes line code
all all
peak peak
0.00B 0.00B 7 def outer():
40.00K 2.00M 8 linear = torch.nn.Linear(100, 100).cuda()
80.00K 2.00M 9 linear2 = torch.nn.Linear(100, 100).cuda()
120.00K 2.00M 10 inner()
## inner
active_bytes reserved_bytes line code
all all
peak peak
80.00K 2.00M 4 def inner():
120.00K 2.00M 5 torch.nn.Linear(100, 100).cuda()
Une explication de ce que signifie chaque colonne peut être trouvée dans la documentation de la torche. Le nom de n'importe quel champ de memory_stats() peut être passé pour display() pour afficher la statistique correspondante.
Si vous utilisez un décorateur profile , les statistiques de mémoire sont collectées pendant plusieurs essais et seul le maximum est affiché à la fin. Nous fournissons également une API plus flexible appelée profile_every qui imprime les informations de mémoire à chaque fois de l'exécution de la fonction. Vous pouvez simplement remplacer @profile par @profile_every(1) pour imprimer l'utilisation de la mémoire pour chaque exécution.
Le @profile et @profile_every peuvent également être mélangés pour prendre plus de contrôle de la granularité de débogage.
class Net ( torch . nn . Module ):
def __init__ ( self ):
super (). __init__ ()
@ profile
def forward ( self , inp ):
#do_somethingset_target_gpu . La sélection du GPU est mondialement, ce qui signifie que vous devez vous rappeler sur quel GPU vous profilez pendant tout le processus: import torch
from pytorch_memlab import profile , set_target_gpu
@ profile
def func ():
net1 = torch . nn . Linear ( 1024 , 1024 ). cuda ( 0 )
set_target_gpu ( 1 )
net2 = torch . nn . Linear ( 1024 , 1024 ). cuda ( 1 )
set_target_gpu ( 0 )
net3 = torch . nn . Linear ( 1024 , 1024 ). cuda ( 0 )
func () Plus d'échantillons peuvent être trouvés dans test/test_line_profiler.py
Assurez-vous que IPython ou avez installé pytorch-memlab avec pip install pytorch-memlab[ipython] .
Tout d'abord, chargez l'extension:
% load_ext pytorch_memlab Cela rend le %mlrun et %%mlrun LIGNE / MAGICS Cell disponible pour une utilisation. Par exemple, dans une nouvelle cellule, exécutez ce qui suit pour profiler une cellule entière
% % mlrun - f func
import torch
from pytorch_memlab import profile , set_target_gpu
def func ():
net1 = torch . nn . Linear ( 1024 , 1024 ). cuda ( 0 )
set_target_gpu ( 1 )
net2 = torch . nn . Linear ( 1024 , 1024 ). cuda ( 1 )
set_target_gpu ( 0 )
net3 = torch . nn . Linear ( 1024 , 1024 ). cuda ( 0 ) Ou vous pouvez invoquer le profileur pour une seule déclaration sur le %mlrun Cell Magic.
import torch
from pytorch_memlab import profile , set_target_gpu
def func ( input_size ):
net1 = torch . nn . Linear ( input_size , 1024 ). cuda ( 0 )
% mlrun - f func func ( 2048 ) Voir %mlrun? pour obtenir de l'aide sur les arguments pris en charge. Vous pouvez définir le périphérique GPU sur le profil, vider les résultats du profilage sur un fichier et renvoyer l'objet LineProfiler pour l'inspection post-profil.
En savoir plus en consultant le cahier de jupyter de démonstration
Comme le profileur de mémoire ne donne que les informations d'utilisation globale de la mémoire par lignes, une information d'utilisation de la mémoire plus bas peut être obtenue par Memory Reporter .
La mémoire du journaliste itère tous les objets Tensor et obtient l'objet UntypedStorage ( Storage ) pour obtenir l'utilisation de la mémoire réelle au lieu du Tensor.size de surface.Size.
Voir UnppedStorage pour des informations détaillées
import torch
from pytorch_memlab import MemReporter
linear = torch . nn . Linear ( 1024 , 1024 ). cuda ()
reporter = MemReporter ()
reporter . report ()Sorties:
Element type Size Used MEM
-------------------------------------------------------------------------------
Storage on cuda:0
Parameter0 (1024, 1024) 4.00M
Parameter1 (1024,) 4.00K
-------------------------------------------------------------------------------
Total Tensors: 1049600 Used Memory: 4.00M
The allocated memory on cuda:0: 4.00M
-------------------------------------------------------------------------------
import torch
from pytorch_memlab import MemReporter
linear = torch . nn . Linear ( 1024 , 1024 ). cuda ()
inp = torch . Tensor ( 512 , 1024 ). cuda ()
# pass in a model to automatically infer the tensor names
reporter = MemReporter ( linear )
out = linear ( inp ). mean ()
print ( '========= before backward =========' )
reporter . report ()
out . backward ()
print ( '========= after backward =========' )
reporter . report ()Sorties:
========= before backward =========
Element type Size Used MEM
-------------------------------------------------------------------------------
Storage on cuda:0
weight (1024, 1024) 4.00M
bias (1024,) 4.00K
Tensor0 (512, 1024) 2.00M
Tensor1 (1,) 512.00B
-------------------------------------------------------------------------------
Total Tensors: 1573889 Used Memory: 6.00M
The allocated memory on cuda:0: 6.00M
-------------------------------------------------------------------------------
========= after backward =========
Element type Size Used MEM
-------------------------------------------------------------------------------
Storage on cuda:0
weight (1024, 1024) 4.00M
weight.grad (1024, 1024) 4.00M
bias (1024,) 4.00K
bias.grad (1024,) 4.00K
Tensor0 (512, 1024) 2.00M
Tensor1 (1,) 512.00B
-------------------------------------------------------------------------------
Total Tensors: 2623489 Used Memory: 10.01M
The allocated memory on cuda:0: 10.01M
-------------------------------------------------------------------------------
import torch
from pytorch_memlab import MemReporter
linear = torch . nn . Linear ( 1024 , 1024 ). cuda ()
linear2 = torch . nn . Linear ( 1024 , 1024 ). cuda ()
linear2 . weight = linear . weight
container = torch . nn . Sequential (
linear , linear2
)
inp = torch . Tensor ( 512 , 1024 ). cuda ()
# pass in a model to automatically infer the tensor names
out = container ( inp ). mean ()
out . backward ()
# verbose shows how storage is shared across multiple Tensors
reporter = MemReporter ( container )
reporter . report ( verbose = True )Sorties:
Element type Size Used MEM
-------------------------------------------------------------------------------
Storage on cuda:0
0.weight (1024, 1024) 4.00M
0.weight.grad (1024, 1024) 4.00M
0.bias (1024,) 4.00K
0.bias.grad (1024,) 4.00K
1.bias (1024,) 4.00K
1.bias.grad (1024,) 4.00K
Tensor0 (512, 1024) 2.00M
Tensor1 (1,) 512.00B
-------------------------------------------------------------------------------
Total Tensors: 2625537 Used Memory: 10.02M
The allocated memory on cuda:0: 10.02M
-------------------------------------------------------------------------------
import torch
from pytorch_memlab import MemReporter
lstm = torch . nn . LSTM ( 1024 , 1024 ). cuda ()
reporter = MemReporter ( lstm )
reporter . report ( verbose = True )
inp = torch . Tensor ( 10 , 10 , 1024 ). cuda ()
out , _ = lstm ( inp )
out . mean (). backward ()
reporter . report ( verbose = True ) Comme indiqué ci-dessous, le (->) indique la réutilisation des mêmes sorties de stockage:
Element type Size Used MEM
-------------------------------------------------------------------------------
Storage on cuda:0
weight_ih_l0 (4096, 1024) 32.03M
weight_hh_l0(->weight_ih_l0) (4096, 1024) 0.00B
bias_ih_l0(->weight_ih_l0) (4096,) 0.00B
bias_hh_l0(->weight_ih_l0) (4096,) 0.00B
Tensor0 (10, 10, 1024) 400.00K
-------------------------------------------------------------------------------
Total Tensors: 8499200 Used Memory: 32.42M
The allocated memory on cuda:0: 32.52M
Memory differs due to the matrix alignment
-------------------------------------------------------------------------------
Element type Size Used MEM
-------------------------------------------------------------------------------
Storage on cuda:0
weight_ih_l0 (4096, 1024) 32.03M
weight_ih_l0.grad (4096, 1024) 32.03M
weight_hh_l0(->weight_ih_l0) (4096, 1024) 0.00B
weight_hh_l0.grad(->weight_ih_l0.grad) (4096, 1024) 0.00B
bias_ih_l0(->weight_ih_l0) (4096,) 0.00B
bias_ih_l0.grad(->weight_ih_l0.grad) (4096,) 0.00B
bias_hh_l0(->weight_ih_l0) (4096,) 0.00B
bias_hh_l0.grad(->weight_ih_l0.grad) (4096,) 0.00B
Tensor0 (10, 10, 1024) 400.00K
Tensor1 (10, 10, 1024) 400.00K
Tensor2 (1, 10, 1024) 40.00K
Tensor3 (1, 10, 1024) 40.00K
-------------------------------------------------------------------------------
Total Tensors: 17018880 Used Memory: 64.92M
The allocated memory on cuda:0: 65.11M
Memory differs due to the matrix alignment
-------------------------------------------------------------------------------
AVIS:
Lors du transfert avec
grad_mode=True, Pytorch maintient les tampons de tenseur pour la rétro-propagation future, au niveau C. Ces tampons ne seront donc pas gérés ou collectés par Pytorch. Mais si vous stockez ces résultats intermédiaires sous forme de variables Python, ils seront signalés.
Vous pouvez également filtrer l'appareil à signaler en passant des arguments supplémentaires: report(device=torch.device(0))
Un exemple défaillant en raison des tampons de tenseur latéral de Pytorch
Dans l'exemple suivant, un tampon temporaire est créé à inp * (inp + 2) pour stocker à la fois inp et inp + 2 , malheureusement Python ne connaît que l'existence d'INP, nous avons donc une mémoire de 2M perdue, qui est de la même taille de Tensor inp .
import torch
from pytorch_memlab import MemReporter
linear = torch . nn . Linear ( 1024 , 1024 ). cuda ()
inp = torch . Tensor ( 512 , 1024 ). cuda ()
# pass in a model to automatically infer the tensor names
reporter = MemReporter ( linear )
out = linear ( inp * ( inp + 2 )). mean ()
reporter . report ()Sorties:
Element type Size Used MEM
-------------------------------------------------------------------------------
Storage on cuda:0
weight (1024, 1024) 4.00M
bias (1024,) 4.00K
Tensor0 (512, 1024) 2.00M
Tensor1 (1,) 512.00B
-------------------------------------------------------------------------------
Total Tensors: 1573889 Used Memory: 6.00M
The allocated memory on cuda:0: 8.00M
Memory differs due to the matrix alignment or invisible gradient buffer tensors
-------------------------------------------------------------------------------
Parfois, les gens souhaitent préempter votre tâche d'exécution, mais vous ne voulez pas enregistrer le point de contrôle, puis charger, en fait, tout ce dont ils ont besoin sont des ressources GPU (généralement des ressources CPU et la mémoire du processeur est toujours épuisée dans les grappes GPU), afin que vous puissiez déplacer tous vos espaces de travail de GPU à CPU, puis arracher votre tâche à partir d'un signal de redémarrage.
Toujours en développement ... mais vous pouvez vous amuser avec:
from pytorch_memlab import Courtesy
iamcourtesy = Courtesy ()
for i in range ( num_iteration ):
if something_happens :
iamcourtesy . yield_memory ()
wait_for_restart_signal ()
iamcourtesy . restore ()Memory_Reporter , les tenseurs intermédiaires ne sont pas couverts correctement, vous pouvez donc insérer de telles logiques de courtoisie après backward ou avant forward .J'ai subi beaucoup de débogage de la mémoire étrange au cours de mes 3 ans de développement de modèles d'apprentissage en profondeur efficaces, et bien sûr, j'ai beaucoup appris de la grande communauté open source.
DataFrame.drop pour Pandas 1.5+ MemReporter (# 24)MemReporter line_profiler introuvable