首先需要确保系统上安装了CUDA支持的NVIDIA GPU和相应的驱动程序。
基本步骤如下
检查CUDA是否可用:
使用 torch.cuda.is_available() 来检查CUDA是否可用。
指定设备:
可以使用 torch.device(“cuda:0”) 来指定要使用的GPU。如果系统有多个GPU,可以通过改变索引(例如 “cuda:1”)来选择不同的GPU。
将模型移至GPU:
使用 .to(device) 方法将模型移至GPU。例如,如果 model 是PyTorch模型,且 device 是之前定义的设备,使用 model.to(device)。
将数据移至GPU:
同样地,将数据(例如张量)也需要移至GPU。例如,如果 data 是一个张量,使用 data = data.to(device)。
进行计算:
然后可以像平常一样进行模型的训练或推断,PyTorch会在GPU上执行计算。
将数据移回CPU(如果需要):
如果你需要将结果移回CPU,可以使用 .to(“cpu”)。
简单的示例代码如下:
import torch
import torch.nn as nn# 检查CUDA是否可用
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")# 创建一个模型示例
model = nn.Linear(10, 5).to(device)# 创建一些数据并移至GPU
input_data = torch.randn(10).to(device)# 在GPU上进行前向传播
output = model(input_data)# 将输出移回CPU(如果需要)
output = output.to("cpu")
在使用GPU进行计算时,要管理好内存,因为GPU内存通常比系统内存有限。如果遇到内存溢出错误,可能需要减少批量大小或优化模型。