如何优化PyTorch中大规模数据集的BERT嵌入提取运行速度？

针对你的问题，以下是不降低嵌入质量的PyTorch专属优化手段，核心围绕批量处理、推理效率提升展开：

1. 批量处理 + PyTorch DataLoader（核心优化） #

单条文本处理的开销（如GPU启动、数据传输）占比极高，批量处理是提升速度的关键。结合PyTorch DataLoader实现自动化批量调度：

from transformers import AutoTokenizer, AutoModel
import torch
from torch.utils.data import Dataset, DataLoader

# 初始化模型和tokenizer，启用快速tokenizer提升编码速度
tokenizer = AutoTokenizer.from_pretrained('bert-base-uncased', use_fast=True)
model = AutoModel.from_pretrained('bert-base-uncased', return_dict=False)
model.eval()  # 切换到评估模式，关闭训练相关层

# 自定义数据集类
class TextDataset(Dataset):
    def __init__(self, texts):
        self.texts = texts
    
    def __len__(self):
        return len(self.texts)
    
    def __getitem__(self, idx):
        return self.texts[idx]

# 批量编码的collate函数
def collate_fn(batch_texts):
    return tokenizer(
        batch_texts,
        return_tensors="pt",
        truncation=True,
        max_length=512,
        padding=True,
        pad_to_multiple_of=8  # 对齐到8的倍数，提升GPU计算效率
    )

# 批量提取嵌入的函数
def extract_batch_embeddings(model, dataloader, device):
    model.to(device)
    all_embeddings = []
    with torch.no_grad():  # 关闭梯度计算，节省内存和计算资源
        for batch in dataloader:
            # 将批量数据移到指定设备
            batch = {k: v.to(device) for k, v in batch.items()}
            outputs = model(**batch)
            # 取<[BOS_never_used_51bce0c785ca2f68081bfa7d91973934]>token的嵌入，形状为(batch_size, hidden_size)
            cls_embeddings = outputs[0][:, 0, :].detach().cpu()
            all_embeddings.append(cls_embeddings)
    # 合并所有批次的嵌入结果
    return torch.cat(all_embeddings, dim=0)

# 示例使用
texts = ["Sample text 1", "Sample text 2", ...]  # 你的大规模文本数据集
dataset = TextDataset(texts)
# 根据GPU内存调整batch_size（如16、32、64，避免OOM）
dataloader = DataLoader(dataset, batch_size=32, shuffle=False, collate_fn=collate_fn)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
embeddings = extract_batch_embeddings(model, dataloader, device)

2. 模型推理效率优化 #

2.1 混合精度推理 #

利用PyTorch的自动混合精度（AMP），在不损失嵌入质量的前提下减少计算量和显存占用：

from torch.cuda.amp import autocast

def extract_batch_embeddings(model, dataloader, device):
    model.to(device)
    all_embeddings = []
    with torch.no_grad(), autocast():  # 启用混合精度
        for batch in dataloader:
            batch = {k: v.to(device) for k, v in batch.items()}
            outputs = model(**batch)
            cls_embeddings = outputs[0][:, 0, :].detach().cpu()
            all_embeddings.append(cls_embeddings)
    return torch.cat(all_embeddings, dim=0)

2.2 模型量化 #

将模型转为INT8精度，大幅提升推理速度，嵌入质量损失可忽略：

• CPU推理推荐动态量化：

model_quantized = torch.quantization.quantize_dynamic(
    model,
    {torch.nn.Linear},
    dtype=torch.qint8
)
# 使用model_quantized替代原模型执行后续推理

• GPU推理推荐静态量化：可通过PyTorch的torch.ao.quantization模块实现，或借助Hugging Face accelerate库简化流程。

2.3 固定模型状态 #

确保始终调用model.eval()，关闭dropout、BatchNorm等训练专属层，避免不必要的计算和随机波动。

3. 细节优化 #

• 调整batch_size：逐步增大batch_size直到出现显存不足（OOM），再回退一个量级，最大化利用GPU显存。
• 数据预加载：若数据集规模适中，可提前完成tokenization并将数据移至GPU，减少每批次的数据传输开销。
• 多进程DataLoader：在CPU预处理阶段启用多进程，设置DataLoader(num_workers=4)（根据CPU核心数调整），加快数据加载速度。

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