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

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

我正在使用GPU运行一些非常大的深度学习模型，当我选择8的批处理大小时，它可以适合内存，但是如果我使用16的批处理大小，则会导致CUDA超出内存错误，我必须终止该过程。

我的问题是，在实际将数据传递到GPU之前，有没有一种方法可以知道数据将在GPU中占据多少？

例如，以下代码是关于如何创建pytorch数据加载器并将每批数据加载器传递给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)

解决方法

我建议在这里使用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），然后您可以乘以批处理大小。