如何计算Bert的内存需求？答案

【问题标题】：How to calculate the memory requirement of Bert?如何计算Bert的内存需求？
【发布时间】：2020-07-24 14:57:47
【问题描述】：

我很好奇transformers.BertModel 的内存使用情况。我想使用预训练模型来转换文本并保存令牌 [CLS] 的输出。没有训练，只有推理。

我对 bert 的输入是 511 个令牌。批量大小为 16 时，我的代码内存不足。 GPU有32GB内存。我的问题是如何估算Bert的内存使用量。

奇怪的是，批处理大小为 32 的另一个作业以相同的设置成功完成。我的代码如下。

# Create dataloader
bs = 16
train_comb = ConcatDataset([train_data, valid_data])
train_dl = DataLoader(train_comb, sampler=RandomSampler(train_data), batch_size=bs)

model = BertModel.from_pretrained('/my_dir/bert_base_uncased/',
                                  output_attentions=False,
                                  output_hidden_states=False)
model.cuda()
out_list = []
model.eval()
with torch.no_grad():
    for d in train_dl:
        d = [i.cuda() for i in d].      # d = [input_ids, attention_mask, token_type_ids, labels] 
        inputs, labels = d[:3], d[3]    # input_ids has shape 16 x 511
        output = model(*inputs)[0][:, 0, :]
        out_list.append(output)

outputs = torch.cat(out_list)

后来我把for循环改成了下面

with torch.no_grad():
    for d in train_dl:
        d = [i.cuda() for i in d[:3]]          # don't care about the labels
        out_list.append(model(*d)[0][:, 0, :]) # remove the intermediary variables
    del d

总而言之，我的问题是：

如何估算Bert的内存使用量？我想用它来估计批量大小。
我的第二个批量大小为 32 的作业成功完成。是因为它有更多的填充物吗？
对提高我的代码中的内存使用效率有什么建议吗？

【问题讨论】：

标签： memory-management out-of-memory huggingface-transformers

【解决方案1】：

经过一番搜索，发现错误是由将输出附加到 GPU 中的列表引起的。使用以下代码，错误消失了。

with torch.no_grad():
    for d in train_dl:
        d = [i.cuda() for i in d[:3]]          
        out_list.append(model(*d)[0][:, 0, :].cpu()) 
    del d

没有.cpu()，内存不断增加

Tensor size: torch.Size([4, 511]), Memory allocated: 418.7685546875MB
Tensor size: torch.Size([4, 768]), Memory allocated: 424.7568359375MB


Tensor size: torch.Size([4, 511]), Memory allocated: 424.7568359375MB
Tensor size: torch.Size([4, 768]), Memory allocated: 430.7451171875MB


Tensor size: torch.Size([4, 511]), Memory allocated: 430.7451171875MB
Tensor size: torch.Size([4, 768]), Memory allocated: 436.7333984375MB

使用 .cpu()，内存不会改变。

Tensor size: torch.Size([128, 511]), Memory allocated: 420.21875MB
Tensor size: torch.Size([128, 768]), Memory allocated: 420.21875MB


Tensor size: torch.Size([128, 511]), Memory allocated: 420.21875MB
Tensor size: torch.Size([128, 768]), Memory allocated: 420.21875MB


Tensor size: torch.Size([128, 511]), Memory allocated: 420.21875MB
Tensor size: torch.Size([128, 768]), Memory allocated: 420.21875MB

【讨论】：