要在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加速进行深度学习模型的训练和推断。