pytorch_memlab
0.3.0
Um laboratório de gerenciamento de memória CUDA simples e preciso para Pytorch, consiste em diferentes partes sobre a memória:
Características:
line_profiler com API simples.%mlrun / %%mlrun / Cell Magic.Índice
pip install pytorch_memlabpip install git+https://github.com/stonesjtu/pytorch_memlabOs erros fora da memória em Pytorch acontecem com frequência, para bordas novas e programadores experientes. Uma razão comum é que a maioria das pessoas realmente não aprende a filosofia de gerenciamento de memória subjacente de Pytorch e GPUs. Eles escreveram códigos de memória em eficiência e reclamaram de Pytorch comer muita memória CUDA.
Neste repositório, vou compartilhar algumas ferramentas úteis para ajudar a depurar a OOM ou inspecionar o mecanismo subjacente, se alguém estiver interessado.
O Profiler de memória é uma modificação do Python's line_profiler , fornece as informações de uso da memória para cada linha de código na função/método especificado.
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 ()Depois que o script termina ou interrompido pelo teclado, ele fornece as seguintes informações de perfil se você estiver em um notebook Jupyter:
Ou as seguintes informações se você estiver em um terminal somente de texto:
## 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()
Uma explicação do que cada coluna significa pode ser encontrada na documentação da tocha. O nome de qualquer campo de memory_stats() pode ser passado para display() para visualizar a estatística correspondente.
Se você usar o decorador profile , as estatísticas de memória serão coletadas durante várias execuções e apenas a máxima é exibida no final. Também fornecemos uma API mais flexível chamada profile_every , que imprime as informações de memória a cada N vezes a execução da função. Você pode simplesmente substituir @profile por @profile_every(1) para imprimir o uso da memória para cada execução.
O @profile e @profile_every também podem ser misturados para ganhar mais controle da granularidade de depuração.
class Net ( torch . nn . Module ):
def __init__ ( self ):
super (). __init__ ()
@ profile
def forward ( self , inp ):
#do_somethingset_target_gpu . A seleção da GPU é globalmente, o que significa que você deve lembrar em qual GPU está perfilando durante todo o processo: 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 () Mais amostras podem ser encontradas em test/test_line_profiler.py
Certifique-se de instalar IPython ou instalou pytorch-memlab com pip install pytorch-memlab[ipython] .
Primeiro, carregue a extensão:
% load_ext pytorch_memlab Isso disponibiliza o %mlrun e %%mlrun Line/Cell Magics para uso. Por exemplo, em uma nova célula, execute o seguinte para o perfil de uma célula inteira
% % 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 você pode invocar o Profiler para obter uma única declaração sobre a Magic Cell %mlrun .
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 ) Veja %mlrun? Para obter ajuda sobre quais argumentos são suportados. Você pode definir o dispositivo GPU como perfil, despejar resultados de perfil em um arquivo e retornar o objeto LineProfiler para inspeção pós-perfil.
Saiba mais ao conferir o caderno de demonstração Jupyter
Como o Profiler de memória fornece apenas as informações gerais de uso da memória por linhas, uma informação de uso de memória mais de baixo nível pode ser obtida pelo Memory Reporter .
O repórter de memória itera todos os objetos Tensor e obtém o objeto UntypedStorage (anteriormente Storage ) subjacente para obter o uso real da memória em vez do Tensor.size de superfície.Size.
Veja UNTYPEDSTORAGE para obter informações detalhadas
import torch
from pytorch_memlab import MemReporter
linear = torch . nn . Linear ( 1024 , 1024 ). cuda ()
reporter = MemReporter ()
reporter . report ()Saídas:
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 ()Saídas:
========= 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 )Saídas:
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 ) Como mostrado abaixo, o (->) indica a reutilização das mesmas saídas de back-end de armazenamento:
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
-------------------------------------------------------------------------------
PERCEBER:
Ao encaminhar com
grad_mode=True, o Pytorch mantém buffers tensores para futura propagação de volta, no nível C. Portanto, esses buffers não serão gerenciados ou coletados por Pytorch. Mas se você armazenar esses resultados intermediários como variáveis Python, elas serão relatadas.
Você também pode filtrar o dispositivo para relatar, passando argumentos extras: report(device=torch.device(0))
Um exemplo fracassado devido aos buffers de tensor C de Pytorch
No exemplo a seguir, um buffer temp é criado em inp * (inp + 2) para armazenar inp e inp + 2 , infelizmente o Python conhece apenas a existência de INP, por isso temos a memória 2M perdida, que é do mesmo tamanho 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 ()Saídas:
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
-------------------------------------------------------------------------------
Às vezes, as pessoas gostariam de antecipar sua tarefa em execução, mas você não deseja economizar e depois carregar, na verdade tudo o que eles precisam são os recursos da GPU (normalmente os recursos da CPU e a memória da CPU são sempre sobressalentes nos clusters de GPU), para que você possa mover todos os seus espaços de trabalho da GPU para a CPU e, em seguida, param e param o sinal de que é um sinal de que é um sinal de que é um sinal de que é um sinal de que você pode suportar.
Ainda se desenvolvendo ..... mas você pode se divertir com:
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 , os tensores intermediários não são cobertos corretamente, portanto, você pode inserir essas lógicas de cortesia após backward ou forward .Sofri muito depurando o uso estranho da memória durante meus três anos de desenvolvimento de modelos de aprendizado profundo eficiente e, é claro, aprendi muito com a grande comunidade de código aberto.
DataFrame.drop para pandas 1.5+ MemReporter (#24)MemReporter line_profiler não encontrado 
