将句子转换器模型从GPU分离到cpu

Question

我在GPU上训练了一个SentenceTransformer模型并保存了它。现在我想在另一台没有GPU的机器上使用它，但是我找不到在cpu上加载它的方法。

from sentence_transformers import SentenceTransformer
model_name = 'all-MiniLM-L6-v2'
model = SentenceTransformer(model_name, device='cuda')

Answer 1

设置device参数为cpu

仅供参考:device以pytorch设备(如cpu, cuda, cuda:0等)为值，默认设置为None，检查是否可以使用GPU。

from sentence_transformers import SentenceTransformer
model_name = 'all-MiniLM-L6-v2'
model = SentenceTransformer(model_name, device='cpu')

Answer 2

您可能需要先安装torch的CPU版本

pip3 install torch --index-url https://download.pytorch.org/whl/cpu

如下所示:https://pytorch.org/get-started/locally/

注意:在这个版本中，也不需要device参数。