pytorch_memlab
0.3.0
Ein einfaches und genaues CUDA -Speichermanagementlabor für Pytorch, das aus verschiedenen Teilen über das Gedächtnis besteht:
Merkmale:
line_profiler -Stil mit einfacher API.%mlrun / %%mlrun Zeilen / Zell Magic -Befehle.Inhaltsverzeichnis
pip install pytorch_memlabpip install git+https://github.com/stonesjtu/pytorch_memlabAußerhalb des Memory-Fehlers in Pytorch treten häufig für Neubien und erfahrene Programmierer auf. Ein häufiger Grund ist, dass die meisten Menschen die zugrunde liegende Gedächtnismanagementphilosophie von Pytorch und GPUs nicht wirklich lernen. Sie schrieben in effizientes Speichercodes und beschwerten sich über Pytorch, der zu viel CUDA-Speicher aß.
In diesem Repo werde ich einige nützliche Tools teilen, um OOM zu debuggen, oder um den zugrunde liegenden Mechanismus zu überprüfen, wenn jemand interessiert ist.
Der Speicherprofiler ist eine Modifikation des Python line_profiler . Er gibt die Speicherverbrauchsinformationen für jede Codezeile in der angegebenen Funktion/Methode an.
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 ()Nachdem das Skript durch die Tastatur fertiggestellt oder unterbrochen wird, gibt es die folgenden Profilerierungsinformationen, wenn Sie sich in einem Jupyter -Notizbuch befinden:
oder die folgenden Informationen, wenn Sie in einem Terminal nur in Text sind:
## 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()
Eine Erläuterung für die gesamte Spalte finden Sie in der Fackeldokumentation. Der Name eines beliebigen Feldes von memory_stats() kann an display() übergeben werden, um die entsprechende Statistik anzuzeigen.
Wenn Sie profile verwenden, werden die Speicherstatistiken während mehrerer Läufe gesammelt und am Ende wird nur die maximale angezeigt. Wir bieten auch eine flexiblere API namens profile_every , die die Speicherinformationen in jeder NAME der Funktionsausführung druckt. Sie können @profile einfach durch @profile_every(1) ersetzen, um die Speicherverwendung für jede Ausführung zu drucken.
Die @profile und @profile_every können auch gemischt werden, um mehr Kontrolle über die Debugging -Granularität zu erlangen.
class Net ( torch . nn . Module ):
def __init__ ( self ):
super (). __init__ ()
@ profile
def forward ( self , inp ):
#do_somethingset_target_gpu umstellen. Die GPU -Auswahl erfolgt weltweit, was bedeutet, dass Sie sich daran erinnern müssen, welche GPU Sie während des gesamten Prozesses profilieren: 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 () Weitere Proben finden Sie in test/test_line_profiler.py
Stellen Sie sicher, dass Sie IPython installiert haben oder pytorch-memlab mit pip install pytorch-memlab[ipython] .
Laden Sie zunächst die Erweiterung:
% load_ext pytorch_memlab Dadurch können die %mlrun und %%mlrun LINE/CELL MAGOCES verwendet werden. Führen Sie beispielsweise in einer neuen Zelle Folgendes aus, um eine ganze Zelle zu profilieren
% % 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 ) Oder Sie können den Profiler für eine einzige Aussage über die %mlrun Cell Magic aufrufen.
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 ) Siehe %mlrun? Für Hilfe bei den Argumenten unterstützt. Sie können das GPU-Gerät auf Profil einstellen, Profilerierungsergebnisse in eine Datei abgeben und das LineProfiler Objekt für die postprofile Inspektion zurückgeben.
Erfahren Sie mehr, indem Sie sich das Demo Jupyter -Notizbuch ansehen
Da der Speicherprofiler nur die allgemeinen Speicherverbrauchsinformationen nach Zeilen angibt, kann von Speicherreporter eine bessere Speicherverwendungsinformation mit niedriger Ebene erhalten werden.
Speicherreporter iteriert alle Tensor und erhält das zugrunde liegende UntypedStorage -Objekt (zuvor Storage ), um die tatsächliche Speicherverwendung anstelle des Tensor.size zu erhalten.
Detaillierte Informationen finden Sie unter UntypedStorage
import torch
from pytorch_memlab import MemReporter
linear = torch . nn . Linear ( 1024 , 1024 ). cuda ()
reporter = MemReporter ()
reporter . report ()Ausgänge:
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 ()Ausgänge:
========= 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 )Ausgänge:
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 ) Wie unten gezeigt, gibt der (->) die Wiederverwendung der gleichen Speicher-Back-End-Ausgänge an:
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
-------------------------------------------------------------------------------
BEACHTEN:
Bei der Weiterleitung mit
grad_mode=Truehält Pytorch Tensorpuffer für die zukünftige Back-Propagation in C-Ebene bei. Diese Puffer werden also nicht von Pytorch verwaltet oder gesammelt. Wenn Sie diese Zwischenergebnisse jedoch als Python -Variablen speichern, werden sie gemeldet.
Sie können das Gerät auch filtern, über das zusätzliche Argumente bestehen: report(device=torch.device(0))
Ein fehlgeschlagenes Beispiel aufgrund von Pytorchs C Side Tensor Puffer
Im folgenden Beispiel wird bei inp * (inp + 2) ein Temperaturpuffer erstellt, um sowohl inp als auch inp + 2 zu speichern. Leider kennt Python nur die Existenz von INP, sodass wir 2 m Speicher verloren haben, was die gleiche Größe des Tensor inp hat.
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 ()Ausgänge:
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
-------------------------------------------------------------------------------
Manchmal möchten die Leute Ihre Laufaufgabe vorbeugen, aber Sie möchten den Kontrollpunkt nicht speichern und dann laden, eigentlich sind sie nur GPU -Ressourcen (in der Regel sind CPU -Ressourcen und CPU -Speicher in GPU -Clustern immer wieder erspart), sodass Sie alle Ihre Arbeitsbereiche von GPU in CPU verschieben können, bis Sie Ihre Aufgabe anstellen, anstatt zu speichern, anstatt das Scheck und das Schütteln zu laden, und das Schaltpunkt und das Boten von Kratzern, das von Kratzern ausgelöst wird.
Immer noch entwickelt ..... aber Sie können Spaß haben mit:
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 erwähnt, werden Zwischenzensoren nicht ordnungsgemäß abgedeckt, sodass Sie nach backward oder vorwärts vorwärts oder vor forward diese doppelte Logik einfügen möchten.Ich habe viel erlitten, wenn ich den seltsamen Gedächtnisgebrauch während meiner dreijährigen Entwicklung effizienter Deep-Learning-Modelle debuguggierte und natürlich viel von der großen Open-Source-Community gelernt habe.
DataFrame.drop für Pandas 1.5+ MemReporter (#24)MemReporter line_profiler nicht gefunden 
