Ubuntu 8GB机器HuggingFace推理OOM：低代码改动的批量处理方案

1. 批量处理输入数据（最直接的低改动方案） #

把整批数据拆分成小批次逐次处理，避免一次性加载所有数据占用内存。基于你现有代码，只需添加少量逻辑实现：

from torch.utils.data import TensorDataset, DataLoader

# 保留原tokenize代码
tokenized_test = tokenizer(dataset["test"]["text"], padding=True, truncation=True, return_tensors="pt")

# 将tokenized结果封装为数据集，按小批次加载
test_dataset = TensorDataset(tokenized_test['input_ids'], tokenized_test['attention_mask'])
# 根据模型大小调整batch_size（比如8/16，内存吃紧就设更小）
test_dataloader = DataLoader(test_dataset, batch_size=8)

logits_list = []
with torch.no_grad():
    for batch in test_dataloader:
        input_ids, attention_mask = batch
        batch_logits = model(input_ids=input_ids, attention_mask=attention_mask).logits
        logits_list.append(batch_logits)

# 合并所有批次的logits结果
logits = torch.cat(logits_list, dim=0)

这种方式完全保留原有核心逻辑，仅通过分批次避免内存过载，数据集再增大也能通过调整batch_size适配。

2. 启用半精度推理（改动极小的内存压缩方案） #

如果你的GPU支持半精度计算（主流NVIDIA/AMD GPU都支持），可以将模型和输入转为半精度浮点数，直接减少约50%的内存占用：

tokenized_test = tokenizer(dataset["test"]["text"], padding=True, truncation=True, return_tensors="pt")

# 模型转半精度
model = model.half()
# 输入张量同步转半精度
tokenized_test = {k: v.half() for k, v in tokenized_test.items()}

with torch.no_grad():
    logits = model(**tokenized_test).logits

或者用自动混合精度（兼容性更好）：

tokenized_test = tokenizer(dataset["test"]["text"], padding=True, truncation=True, return_tensors="pt")

with torch.no_grad(), torch.cuda.amp.autocast():
    logits = model(**tokenized_test).logits

仅需添加1-2行代码，无需修改数据处理逻辑，适合快速缓解内存压力。

3. 批量+半精度组合（最优适配方案） #

如果单一种方式仍无法解决OOM，可将两者结合，进一步压缩内存占用：

from torch.utils.data import TensorDataset, DataLoader

tokenized_test = tokenizer(dataset["test"]["text"], padding=True, truncation=True, return_tensors="pt")
# 输入转半精度
tokenized_test = {k: v.half() for k, v in tokenized_test.items()}

test_dataset = TensorDataset(tokenized_test['input_ids'], tokenized_test['attention_mask'])
test_dataloader = DataLoader(test_dataset, batch_size=16) # 半精度下可适当调大batch_size

logits_list = []
model = model.half()
with torch.no_grad():
    for batch in test_dataloader:
        input_ids, attention_mask = batch
        batch_logits = model(input_ids=input_ids, attention_mask=attention_mask).logits
        logits_list.append(batch_logits)

logits = torch.cat(logits_list, dim=0)

额外小技巧 #

• 若用CPU推理，可将tokenizer的return_tensors设为"np"，用numpy数组分批次后再转torch张量，降低CPU内存占用。
• 避免在循环中创建冗余张量，尽量复用变量。

内容的提问来源于stack exchange，提问作者Lucas Azevedo