Pytorch使用GPU

import torch
from torch import nn

判断是否可用GPU

torch.cuda.is_available()

True 可用 False不可用

torch.cuda.device_count() #查询有几个GPU可用

定义使用GPU和CPU的函数

def try_gpu(i=0):  #@save
    """如果存在，则返回gpu(i)，否则返回cpu()"""
    if torch.cuda.device_count() >= i + 1:
        return torch.device(f'cuda:{i}')
    return torch.device('cpu')

def try_all_gpus():  #@save
    """返回所有可用的GPU，如果没有GPU，则返回[cpu(),]"""
    devices = [torch.device(f'cuda:{i}')
             for i in range(torch.cuda.device_count())]
    return devices if devices else [torch.device('cpu')]

try_gpu(), try_gpu(10), try_all_gpus()

神经网络在GPU上运行

net = nn.Sequential(nn.Linear(3,1))
net = net.to(device = try_gpu())  # net.to() 选择gpu
net(X)

net[0].weight.data.device  #查看模型参数存贮在那个位置