PyTorch——使用GPU

参考链接

  1. https://tangshusen.me/Dive-into-DL-PyTorch/#/chapter04_DL_computation/4.6_use-gpu

查看GPU信息

!nvidia-smi  # 对Linux/macOS用户有效

查看GPU是否可用

torch.cuda.is_available()

查看GPU数量

torch.cuda.device_count()

查看当前GPU索引号

torch.cuda.current_device()

查看GPU名字

torch.cuda.get_device_name(0)

Tensor的GPU计算

使用.cuda()可以将CPU上的Tensor转换(复制)到GPU上。如果有多块GPU,我们用.cuda(i)来表示第 ii 块GPU及相应的显存(ii从0开始)且cuda(0)和cuda()等价。

x = x.cuda(0)

通过Tensor的device属性来查看该Tensor所在的设备。

x.device

我们可以直接在创建的时候就指定设备。

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

x = torch.tensor([1, 2, 3], device=device)
# or
x = torch.tensor([1, 2, 3]).to(device)

模型的GPU计算

同Tensor类似,PyTorch模型也可以通过.cuda()转换到GPU上。我们可以通过检查模型的参数的device属性来查看存放模型的设备。

net = nn.Linear(3, 1)
print(list(net.parameters())[0].device)

net.cuda()
print(list(net.parameters())[0].device)
已标记关键词 清除标记
相关推荐
©️2020 CSDN 皮肤主题: 数字20 设计师:CSDN官方博客 返回首页