How to enable GPU for SetFit?

Question

I am following this tutorial for SetFit: https://www.philschmid.de/getting-started-setfit

When the training is running, it is using my CPU instead of my GPU. Is there a way I can enable it?

Here is the main part of the code:

from setfit import SetFitModel, SetFitTrainer
from sentence_transformers.losses import CosineSimilarityLoss

# Load a SetFit model from Hub
model_id = "sentence-transformers/all-mpnet-base-v2"
model = SetFitModel.from_pretrained(model_id)

# Create trainer
trainer = SetFitTrainer(
    model=model,
    train_dataset=train_dataset,
    eval_dataset=test_dataset,
    loss_class=CosineSimilarityLoss,
    metric="accuracy",
    batch_size=64,
    num_iterations=20, # The number of text pairs to generate for contrastive learning
    num_epochs=1, # The number of epochs to use for constrastive learning
)

# Train and evaluate
trainer.train()
metrics = trainer.evaluate()

Answer 1

If your training is running on CPU rather than GPU, it is because:

1. Either you installed the CPU version of the PyTorch.
2. Either the version of CUDA/CUDNN and PyTorch are not compatible, and the training falls back to CPU instead of GPU.

In essence it has nothing to do with the SetFit model.

A working example for me in recent projects is:

(1) pip/pip3 install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu116

(2) pip install transformers==4.22.0

Note that you may have to uninstall pytorch first before reinstalling it: pip uninstall pytorch.

In order to make sure your GPU is visible, a short print would suffice:

training_device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")