查看N卡信息

nvidia-smi

Pytorch中,每个数组都有1个device,通常称之为context,默认情况下,所有的变量和相关的计算都分配给CPU。 多GPU 多卡计算一定是大趋势

计算设备

指定用于存储和计算的设备

import torch
from torch import nn

torch.device('cpu'),torch.device('cuda'),torch.device('cuda:1')

查询可用gpu的数量

torch.cuda.device_count()

定义2个方便的函数,允许我们在不存在所有GPU的情况下运行代码

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}') return devices if devices else [torch.device('cpu')]]

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

查询张量所在的设备,默认情况下是在cpu上创建的

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

存储在GPU上

可以在创建张量时候指定存储设备,接下来在第一个gpu上创建张量变量X,在GPU上创建的张量只消耗这个GPU的显存。

X = torch.ones(2,3,device=try_gpu())
X

在第2个GPU上创建1个随机张量

Y = torch.rand(2,3,device=try_gpu())
Y

复制

计算X+Y X在gpu(0)上,Y在gpu(1)上,需要将X移到gpu(1),才能完成相加计算

Z = X.cuda(1)
print(X)
print(Z)
Y+Z

要是调用Z.cuda(1)会怎样,还是会返回Z,而不是复制并分配新内存

Z.cuda(1) is Z

神经网络与GPU

神经网络模型也可以指定gpu

net = nn.Sequential(nn.Linear(3,1))
net = net.to(device=try_gpu())
net(X)

模型参数存储在同一个GPU上

net[0].weight.data.device