【动手学习pytorch笔记】8.GPU相关

GPU相关

查看计算机有多少台GPU

import torch 
from torch import nn

torch.cuda.device_count()

输出

1

定义一些常用函数

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()

try_gpu()有就返回GPU，没有就返回CPU

try_all_gpus()返回所有GPU

输出

(device(type='cuda', index=0),
 device(type='cpu'),
 [device(type='cuda', index=0)])

在GPU上使用张量

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

输出

device(type='cpu')

默认是在CPU上运行的

指定在GPU上运行

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

输出

tensor([[1., 1., 1.],
        [1., 1., 1.]], device='cuda:0')

张量的计算

x + X

---

RuntimeError                              Traceback (most recent call last)

<ipython-input-20-cade78b2014d> in <module>
----> 1 x + X

RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!

可以看到不再一起是不能做运算的

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

输出

tensor([[0.8074, 0.9579, 0.0942],
        [0.1074, 0.5444, 0.1546]], device='cuda:0')

X + Y

输出

tensor([[1.8074, 1.9579, 1.0942],
        [1.1074, 1.5444, 1.1546]], device='cuda:0')

在GPU上使用神经网络

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

net(X)

输出

tensor([[0.0069],
        [0.0069]], device='cuda:0', grad_fn=<AddmmBackward0>)

查看参数位置

net[0].weight.data.device

输出

device(type='cuda', index=0)