PyTorch实现的BERT嵌入提取函数内存占用过快过高问题求助

核心问题：未关闭自动微分导致计算图冗余存储 #

你当前的代码运行在PyTorch默认的自动微分模式下，存在以下未纳入测算的内存开销：

• 推理过程不需要计算梯度，但PyTorch会默认保留所有前向传播的中间层张量、计算图节点用于后续反向传播，仅BERT一次前向传播的中间张量体积就是最终输出last_hidden_state的10~20倍，这部分开销完全没有被你纳入内存测算范围，是内存暴涨的核心原因。
• 所有产出的embeddings、拼接得到的chunk_embeddings都会携带梯度计算引用，即使函数返回后，只要返回的val还存在，整个迭代过程中产生的所有计算图、中间张量都不会被垃圾回收，自然会出现“返回值很小但内存占用降不下来”的现象。

次要问题：反复拼接张量的冗余内存开销 #

你使用torch.cat循环拼接空张量的写法，每次迭代都会生成一份新的完整chunk_embeddings副本，旧的副本在未被回收的情况下会产生双倍内存占用，迭代次数越多冗余开销越大。

存在的认知错误 #

你之前的内存测算仅统计了最终产出张量的纯数据体积，完全忽略了PyTorch深度学习框架运行时的额外开销：自动微分的计算图、中间张量、张量对象的元数据、张量拼接的临时副本等，这部分开销远大于你测算的纯数据体积。

def get_chunk_embeddings(encoded_dataset, batch_size):
    # 提前预分配输出张量，避免反复拼接的冗余开销
    total_samples = len(encoded_dataset['input_ids'])
    chunk_embeddings = torch.empty((total_samples, 768), device='cpu') # 直接存在CPU内存更省空间
    # 关闭自动微分，推理场景下完全不生成计算图，中间张量用完即回收
    with torch.no_grad():
        # 向上取整计算迭代次数，兼容样本总数不能被batch_size整除的情况
        for i in range((total_samples + batch_size - 1) // batch_size):
            start_idx = i * batch_size
            end_idx = min(start_idx + batch_size, total_samples)
            input_ids = encoded_dataset['input_ids'][start_idx:end_idx]
            attention_mask = encoded_dataset['attention_mask'][start_idx:end_idx]
            # 优先用模型__call__方法而非直接调用forward，兼容钩子等内置逻辑
            outputs = model(input_ids=input_ids, attention_mask=attention_mask)
            embeddings = outputs['last_hidden_state']
            # 带掩码的均值计算
            mask = attention_mask.unsqueeze(-1)
            embeddings = (embeddings * mask).sum(dim=1) / attention_mask.sum(dim=1, keepdim=True)
            # 直接写入预分配的张量，无需拼接
            chunk_embeddings[start_idx:end_idx] = embeddings.cpu()
    return chunk_embeddings

额外优化建议 #

• 如果使用GPU运行模型，拿到embeddings后立刻调用.cpu()转移到CPU内存，避免GPU显存和CPU内存中同时留存冗余副本。
• 如果不需要保留梯度，返回的结果可以直接调用.detach()彻底切断梯度引用，进一步释放冗余内存。

内容的提问来源于stack exchange，提问作者Luca Guarro