为了在使用 Sentence Transformers 生成嵌入时提高吞吐量，批量处理至关重要。 不是一次处理一个句子，而是将它们分组到批次中并一起处理。 这种方法最大限度地提高了硬件效率，尤其是在 GPU 上，GPU 擅长并行计算。 通过减少重复模型调用的开销并利用并行处理，您可以显着加快大型数据集的嵌入生成速度。
主要方法是使用 Sentence Transformers 库提供的 encode 函数，该函数支持开箱即用的批处理。 例如，如果您有一个包含 1,000 个句子的列表，您可以将整个列表传递给 model.encode()，并指定 batch_size 参数（例如，batch_size=64）。 该库会自动将列表拆分为批次并按顺序处理它们。 较大的批量大小通常会提高吞吐量，但您必须平衡这一点与 GPU 内存限制。 例如，根据模型和硬件能力，批量大小为 64 可能在大致相同的时间内处理 64 个句子，就像单个句子一样。 测试不同的批量大小是找到最佳平衡的关键。
一些优化可以进一步提高性能。 首先，按长度对句子进行预排序可以最大限度地减少批次内的填充，因为 tokenizer 会将序列填充到批次中最长的句子。 排序确保长度相似的句子分组在一起，从而减少浪费的计算。 其次，如果您的 GPU 支持，使用混合精度推理 (FP16) 可以减少内存使用并加快处理速度。 第三，通过设置 convert_to_tensor=True 将数据保存在 GPU 上，可以避免 CPU 和 GPU 之间昂贵的传输。 例如，model.encode(sentences, batch_size=128, convert_to_tensor=True, device='cuda') 直接在 GPU 上处理批次。 最后，对于极其大的数据集，将数据分成块并以适当的批量大小按顺序处理它们可以防止内存不足错误，同时保持吞吐量。