为什么 Tensorflow 报告 CUDA 内存不足但 empty_cache 不起作用？

Question

device = torch.device("cuda:0")
model = BertModel.from_pretrained("bert-base-uncased", output_hidden_states = True)

model.to(device)

train_hidden_states = []
model.eval()

for batch in train_dataloader:
    b_input_ids = batch[0].to(device)
    b_input_mask = batch[1].to(device)

    with torch.no_grad():        
        output = model(b_input_ids, 
                        token_type_ids=None, 
                        attention_mask=b_input_mask,
                        )
        hidden_states = output[2][12]
        train_hidden_states.append(hidden_states)

这里我试图获取train_dataloader中数据的Bert模型的最后一层嵌入。

问题是 CUDA 在 14 个批次后内存不足。

我尝试清空缓存，但它只降低了一点点 GPU 使用率。

with torch.cuda.device('cuda:0'):
    torch.cuda.empty_cache()

可能是什么问题？

Answer 1

您正在 train_hidden_states 列表中的 GPU 上存储张量。您可以在推送到列表 train_hidden_states.append(hidden_states.cpu()).

之前移动到 CPU

为什么 Tensorflow 报告 CUDA 内存不足但 empty_cache 不起作用？

Why Tensorflow reports CUDA out of memory but empty_cache doesn't work?

python

gpu

pytorch