Como saber se os dados passados para a GPU causarão falta de memória no CUDA ou não
Estou usando GPU para executar alguns modelos de aprendizado profundo muito grandes, quando escolho um tamanho de lote de 8, ele pode caber na memória, mas se eu usar um tamanho de lote de 16, isso causará um erro de falta de memória CUDA, e eu tenho que matar o processo.
Minha pergunta é, antes de realmente passar os dados para a GPU, existe uma maneira de saber o tamanho que os dados ocuparão na GPU?
Por exemplo, o código a seguir é sobre como eu crio um dataloader pytorch e passo cada lote do dataloader para a GPU, posso saber o tamanho dele antes de chamarbatch.to(device)
train_dataloader = DataLoader(train_data, sampler=train_sampler, batch_size=batch_size)
for step, batch in enumerate(train_dataloader):
b_input_ids = batch[0].to(device)
b_input_mask = batch[1].to(device)
b_labels = batch[2].to(device)
Respostas
Eu recomendaria usar o pacote archsummary aqui.
pip install torchsummary
e em uso
from torchsummary import summary
myModel.cuda()
summary(myModel, (shapeOfInput)) # where shapeOfInput is a tuple of the sample's dimensions
Isso fornecerá o tamanho do modelo, o tamanho da passagem direta e o tamanho da passagem reversa em MB para um tamanho de lote de 1 e você poderá multiplicar pelo tamanho do lote.