1.查看GPU状态
!nvidia-smi
2.使用计算的设备
torch.device('cpu'),torch.cuda.device('cuda')#这里默认为cuda:0
torch.device('cuda:1')
3.查看可用GPU数量
torch.cuda.device_count()
4. 自定义访问GPU的函数try_gpu(i) 和访问所有GPU的函数try_all_gpus()
def try_gpu(i=0):
if torch.cuda.device_count()>=i+1:
return torch.device(f'cuda:{i}')
return torch.device('cpu')
def try_all_gpus():
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()
返回:(device(type='cuda', index=0), device(type='cpu'), [device(type='cuda', index=0)])
5.查看张量所在的设备:
x = torch.tensor([1,2,3])
x.device
6.将张量存储在GPU上
X = torch.ones(2,3,device=try_gpu())
#这里try_gpu是前面定义好的函数,表示在第1个GPU上创建tensor,核心是torch.device("cuda:1")
X
Y = torch.rand(2,3,device = try_gpu(1))
结果:
tensor([[1., 1., 1.], [1., 1., 1.]], device='cuda:0')
7.计算位置:如果想让X+Y在GPU上计算,需要这两个张量保存在一个GPU上,所以首先需要将他们移动到一个GPU上面,注意:在CPU和GPU之间移动数据会降低性能,应该尽量避免
Z = X.cuda(1)#将X的数据移动到第二个GPU上面
Y+Z#此时在第二个GPU上运算
8.使用GPU进行神经网络计算:net.to(device=try_gpu())函数可以将网络挪到GPU设备上
net = nn.Sequential(nn.Linear(3,1))
net = net.to(device=try_gpu())
net(X)
9.确认模型参数存储在GPU上
net[0].weight.data.device