machine-learning - 如何知道传递给GPU的数据是否会导致CUDA内存不足

Question

我正在使用GPU来运行一些非常大的深度学习模型，当我选择batch size为8时，它可以适合内存，但是如果我使用batch size为16，则会导致CUDA out-of-内存错误，我必须终止该进程。

我的问题是，在实际将数据传递到GPU之前，有没有办法知道数据在GPU中占用了多大？

例如，下面的代码是关于我如何创建一个 pytorch dataloader 并将每批 dataloader 传递到 GPU，我可以在调用 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)

Answer 1

我建议在这里使用 torchsummary 包。

pip install torchsummary

正在使用

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

这将为您提供模型的大小、前向传递的大小和后向传播的大小(以 MB 为单位)，批处理大小为 1，然后您可以乘以批处理大小。