Python pytorch函数消耗内存过快

Question

我正在用 pytorch 编写一个函数，该函数通过变换器模型提供输入，然后通过计算沿特定轴的平均值（使用掩码定义的索引子集）来压缩最后一个嵌入层。由于模型的输出非常非常大，我需要批量处理输入。

我的问题与此函数的逻辑无关，因为我相信我有正确的实现。我的问题是我编写的函数消耗内存的速度过快，几乎无法使用。

这是我的函数：

def get_chunk_embeddings(encoded_dataset, batch_size):
  chunk_embeddings = torch.empty([0,768])
  for i in range(len(encoded_dataset['input_ids'])//batch_size):
    input_ids = encoded_dataset['input_ids'][i*batch_size:i*batch_size + batch_size]
    attention_mask = encoded_dataset['attention_mask'][i*batch_size:i*batch_size + batch_size]
    embeddings = model.forward(input_ids=input_ids, attention_mask=attention_mask)['last_hidden_state']
    embeddings = embeddings * attention_mask[:,:,None]
    embeddings = embeddings.sum(dim=1)/attention_mask.sum(dim=1)[:,None]
    chunk_embeddings = torch.cat([chunk_embeddings, embeddings],0)
  return chunk_embeddings

现在让我们谈谈记忆（下面的数字假设我通过了 batch_size of 8）：

我正在使用 google colab，我有大约 25 GB 的可用内存
model 是一个 BERT 模型，占用 413 MB
encoded_dataset 消耗 0.48 GB
input_ids 消耗 0.413 MB
attention_mask 消耗 4.096 KB
embeddings 在其峰值消耗时消耗 12.6 MB
chunk_embeddings 每次迭代增加 0.024576 MB

所以根据我的理解，我应该能够让 chunk_embeddings 增长到：25GB - 413MB - 0.48GB - 0.413MB - 4.096KB - 12.6MB ~= 24 GB。足够近 100 万次迭代..

在这里，我将通过一个我正在经历的例子来说明：

在运行启用我的函数之前，google colab 告诉我我有足够的内存

现在，为了举例，我将运行函数（仅进行 3 次迭代）明确地说，我把它放在我的 for 循环的末尾： if (i == 2):return chunk_embeddings
现在我运行代码val = get_chunk_embeddings(train_encoded_dataset, 8) 因此，即使只有 3 次迭代，我也以某种方式消耗了将近 5.5 GB 的 RAM。

为什么会这样？同样，在我从函数返回后，所有局部变量都应该被删除，而且 val 不可能这么大。

谁能告诉我我做错了什么或不明白？如果需要更多信息，请告诉我。

Answer 1

为了扩展@GoodDeeds 的答案，默认情况下 pytorch.nn 模块（模型）中的计算会创建计算图并保留梯度（除非您使用 with torch.no_grad() 或类似的东西。这意味着在循环的每次迭代中，嵌入的计算图存储在张量 embeddings 中。embeddings.grad 可能比 embeddings 本身大得多，因为每个层值的梯度与保持对每个前一层值的尊重。接下来，由于您使用 torch.cat，因此您将 embeddingsd 和相关的梯度附加到 chunk_embeddings。这意味着经过几次迭代后，chunk_embeddings 存储了海量的梯度值，这就是你的记忆力所在，有几个解决方案：

如果您需要使用块嵌入进行反向传播（即训练），您应该将损失计算和优化器步骤移动到循环中，以便之后自动清除梯度。
如果仅在推理期间使用此函数，您可以使用 torch.no_grad() 完全禁用梯度计算（这也应该稍微加快计算速度），或者您可以使用 torch.detach() embeddings 在评论中建议的每次迭代中。

示例：

def get_chunk_embeddings(encoded_dataset, batch_size):
  with torch.no_grad():
    chunk_embeddings = torch.empty([0,768])
    for i in range(len(encoded_dataset['input_ids'])//batch_size):
      input_ids = encoded_dataset['input_ids'][i*batch_size:i*batch_size + batch_size]
      attention_mask = encoded_dataset['attention_mask'][i*batch_size:i*batch_size + batch_size]
      embeddings = model.forward(input_ids=input_ids, attention_mask=attention_mask)['last_hidden_state']
      embeddings = embeddings * attention_mask[:,:,None]
      embeddings = embeddings.sum(dim=1)/attention_mask.sum(dim=1)[:,None]
      chunk_embeddings = torch.cat([chunk_embeddings, embeddings],0)
  return chunk_embeddings

Python pytorch函数消耗内存过快

Python pytorch function consumes memory excessively quickly

python

ram

memory-leaks

memory-management

pytorch