Como saber se os dados passados ​​para a GPU causarão falta de memória no CUDA ou não

Aug 17 2020

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

3 David Aug 17 2020 at 15:36

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.