如何使用 GPU 加速 Sentence Transformers 的嵌入生成？代码中需要进行哪些更改？

要使用 GPU 加速 Sentence Transformers 的嵌入生成，你需要确保模型和数据已通过 PyTorch 的 CUDA 支持转移到 GPU。 Sentence Transformers 构建于 PyTorch 之上，如果指定了设备，它会自动利用 GPU 加速。 主要的代码更改包括将模型设置为在初始化期间使用 cuda 而不是 cpu。 例如，加载模型时，你可以将 device='cuda' 传递给 SentenceTransformer 构造函数。 如果模型已加载到 CPU 上，你可以使用 model.to('cuda') 将其移动到 GPU。 如果模型已经配置为 CUDA，则输入数据（例如文本句子）将在嵌入生成期间自动在 GPU 上处理。

一个关键的考虑因素是确保你的输入与 GPU 处理兼容。 当使用 encode() 方法生成嵌入时，Sentence Transformers 会在内部将文本输入转换为张量，并将它们移动到与模型相同的设备。 例如，如果模型位于 GPU 上，则 embeddings = model.encode(sentences) 将自动处理设备放置。 但是，如果你正在处理大型数据集，则批处理对于最大化 GPU 利用率至关重要。 你可以在 encode() 中将 batch_size 参数设置为更高的值（例如，32 或 64）以并行化计算。 较大的批次利用了 GPU 的并行性，但必须保持在内存限制范围内——监控 VRAM 使用情况以避免内存不足错误。

此外，显式设备管理可以提高控制能力。 例如，你可以使用 torch.cuda.is_available() 检查 GPU 是否可用，并有条件地设置设备。 这是一个代码片段

from sentence_transformers import SentenceTransformer
import torch

device = 'cuda' if torch.cuda.is_available() else 'cpu'
model = SentenceTransformer('all-MiniLM-L6-v2', device=device)
sentences = ["This is an example sentence.", "Another example."]
embeddings = model.encode(sentences, batch_size=32, convert_to_tensor=True)

设置 convert_to_tensor=True 可以将嵌入保留在 GPU 上，这对于下游的 GPU 加速任务非常有用。 如果你需要 numpy 数组，请省略此参数，或者在使用 .cpu() 转换之前将张量移回 CPU。 无需其他代码更改——该库会处理底层操作，如数据传输和内核优化。