如何为SetFit启用GPU加速训练？

先确认环境支持CUDA #

首先运行以下代码验证PyTorch是否能识别GPU：

import torch
print(torch.cuda.is_available())

如果输出True，说明环境具备GPU加速条件；如果是False，需要安装对应CUDA版本的PyTorch（根据你的显卡CUDA版本选择适配的安装包）。

修改代码启用GPU #

针对你的核心代码，只需做以下调整：

1. 加载模型时指定GPU设备
   在SetFitModel.from_pretrained中添加device="cuda"参数，让模型直接加载到GPU：

model = SetFitModel.from_pretrained(model_id, device="cuda")

2. （可选）强制训练器绑定GPU
   如果上述设置后仍未生效，可以在初始化SetFitTrainer时显式指定设备：

trainer = SetFitTrainer(
    model=model,
    train_dataset=train_dataset,
    eval_dataset=test_dataset,
    loss_class=CosineSimilarityLoss,
    metric="accuracy",
    batch_size=64,
    num_iterations=20,
    num_epochs=1,
    device="cuda"  # 添加该行绑定GPU
)

验证GPU运行状态 #

训练时可以通过nvidia-smi命令（Linux/macOS）或任务管理器（Windows）查看GPU负载，也可以在代码中打印模型所在设备确认：

print(model.device)

正常输出应为cuda:0（或对应GPU设备标识）。

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