Pytorch&GPU的简单使用

# 查看显卡信息
!nvidia-smi

import torch
from torch import nn

torch.device('cpu')
# 可以查看GPU的数量
torch.cuda.device_count()

下面这两个函数可以试探有没有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_gpu():
    devices = [torch.device(f"cuda:{i}") for i in range(torch.cuda.device_count())]
    return devices if devices else [torch.device('cpu')]

try_gpu(0)
try_all_gpu()

# 查询张量所在设备
x = torch.tensor([1,2,3,4])
x.device

# 将张量创建在GPU上面
X = torch.randn((3,4),device=try_gpu(0))
Y = torch.ones((3,4),device=try_gpu(0))

print(X,Y)

# 在GPU上计算
X = X.cuda(0) # 在GPU之间进行复制，这里由于只有一块GPU所以，相当于原地跳了一下
Z = X+Y

print(Z)

# 怎么样在GPU上计算神经网络呢
net = nn.Sequential(nn.Linear(3,4))
net = net.to(device=try_gpu())

x = torch.randn(9,device=try_gpu()).reshape((3,3))
net(x)
net[0].weight.data.device