【Pytorch写法】tensor与GPU

提示：文章写完后，目录可以自动生成，如何生成可参考右边的帮助文档

---

前言

参考资料
李沐17使用和购买GPU
d2l

正文

查看GPU

!nvidia-smi

查看GPU数量

torch.cuda.device_count()

使用GPU

import torch
from torch import nn

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

单卡/多卡使用GPU

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

张量

x = torch.tensor([1, 2, 3])
x.device # device(type='cpu')

张量在GPU上

X = torch.ones(2, 3, device=try_gpu()) # tensor([[1., 1., 1.],[1., 1., 1.]], device='cuda:0')
Y = torch.rand(2, 3, device=try_gpu(1)) # tensor([[0.1206, 0.2283, 0.4548],[0.9806, 0.9616, 0.0501]], device='cuda:1')

跨卡计算

Z = X.cuda(1) # X device='cuda:0' Z device='cuda:1'
Y + Z # Y device='cuda:1'
Z.cuda(1) is Z # True

网络放在GPU上

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

确认参数在同一个GPU上

net[0].weight.data.device