要在PyTorch中使用GPU加速,需要遵循以下步驟:
檢查是否安裝了CUDA驅動程序和cuDNN庫。這些是GPU加速所需的必要組件。
安裝PyTorch GPU版本。可以通過以下命令安裝:
pip install torch torchvision
確保安裝的是GPU版本而不是CPU版本。
import torch
print(torch.cuda.is_available())
如果輸出為True,則表示PyTorch可以訪問GPU。
# 將模型移動到GPU
model = model.to('cuda')
# 將數據移到GPU
inputs = inputs.to('cuda')
labels = labels.to('cuda')
outputs = model(inputs)
loss = loss_function(outputs, labels)
# 反向傳播
loss.backward()
torch.cuda.empty_cache()
通過以上步驟,就可以在PyTorch中使用GPU加速進行深度學習模型的訓練和推斷。