多次推理时PyTorch出现CUDA内存不足错误求助

针对你在使用all-distilroberta-v1处理pubmed训练集时遇到的显存溢出问题，以下是几个直接有效的解决方法：

1. 分批次处理单文档内的句子 #

你的代码当前会把单个文档的所有句子一次性送入模型，若遇到句子数量极多的文档（比如几十上百句），会瞬间占用大量显存。改成小批次处理单文档的句子，能显著降低峰值显存占用：

def inference(document, batch_size=8):
    sentences = nltk.sent_tokenize(document)
    # 拆分句子为多个小批次
    for i in range(0, len(sentences), batch_size):
        batch_sentences = sentences[i:i+batch_size]
        tokenized_sentences = tokenizer(
            batch_sentences, 
            padding=True, 
            truncation=True, 
            return_tensors="pt"
        ).to('cuda')
        with torch.no_grad():
            model(**tokenized_sentences)
        # 手动释放当前批次的显存
        del tokenized_sentences
        torch.cuda.empty_cache()

可以根据你的显存大小调整batch_size（比如4、8、16），确保单批次显存占用在安全范围内。

2. 启用半精度（FP16）推理 #

将模型和输入转为FP16格式，能直接把显存占用降低约50%，且对推理效果影响极小：

# 加载模型时启用FP16
model = transformers.AutoModel.from_pretrained(
    "sentence-transformers/all-distilroberta-v1",
    torch_dtype=torch.float16
).cuda()

# 推理时也用FP16
def inference(document, batch_size=8):
    sentences = nltk.sent_tokenize(document)
    for i in range(0, len(sentences), batch_size):
        batch_sentences = sentences[i:i+batch_size]
        tokenized_sentences = tokenizer(
            batch_sentences, 
            padding=True, 
            truncation=True, 
            return_tensors="pt"
        ).to('cuda', dtype=torch.float16)  # 转为FP16
        with torch.no_grad():
            model(**tokenized_sentences)
        del tokenized_sentences
        torch.cuda.empty_cache()

3. 启用梯度检查点减少模型显存占用 #

梯度检查点会牺牲少量速度来降低模型的显存占用，适合显存紧张的场景：

model = transformers.AutoModel.from_pretrained(
    "sentence-transformers/all-distilroberta-v1"
).cuda()
# 启用梯度检查点
model.gradient_checkpointing_enable()

4. 调整PyTorch内存分配策略解决碎片问题 #

按错误提示设置环境变量，优化CUDA内存分配，减少碎片：

在运行代码前设置环境变量（Linux/macOS用终端，Windows用命令提示符）：

export PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:128

或者在Python代码开头添加：

import os
os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'max_split_size_mb:128'

这个设置让PyTorch在分配内存时更倾向于拆分出小块内存，减少大内存块分配失败的概率。

内容的提问来源于stack exchange，提问作者Aldan Creo