How to know whether the data passed to GPU will cause CUDA out of memory or not

user_10631590 0
Viewed 238

I am using GPU to run some very large deep learning models, when I choose a batch size of 8, it can fit into the memory, but if I use a batch size of 16, it will cause CUDA out-of-memory error, and I have to kill the process.

My question is, before actually passing the data into GPU, is there a way that I could know how large the data will occupy in the GPU?

For example, the following code is about how I create a pytorch dataloader and pass each batch of the dataloader to the GPU, could I know how large it is before I call batch.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)
pytorch deep-learning machine-learning gpu
1 Answers

I would recommend using the torchsummary package here.

pip install torchsummary

and in use

from torchsummary import summary
myModel.cuda()
summary(myModel, (shapeOfInput)) # where shapeOfInput is a tuple of the sample's dimensions

This will give you the size of the model, the size of the forward pass, and the size of the backpass in MB for a batch size of 1, and you can then multiple out by your batch size.

user_66836390
Related