!nvidia-smi

Mon Aug 16 14:23:34 2021       
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 430.26       Driver Version: 430.26       CUDA Version: 10.2     |
|-------------------------------+----------------------+----------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|===============================+======================+======================|
|   0  GeForce RTX 208...  Off  | 00000000:01:00.0 Off |                  N/A |
| 22%   32C    P8     1W / 250W |   9475MiB / 11018MiB |      0%      Default |
+-------------------------------+----------------------+----------------------+
                                                                               
+-----------------------------------------------------------------------------+
| Processes:                                                       GPU Memory |
|  GPU       PID   Type   Process name                             Usage      |
|=============================================================================|
|    0      7024      C   ...t/anaconda3/envs/pytorch0.4/bin/python3  9465MiB |
+-----------------------------------------------------------------------------+

import torch
from torch import nn

# torch.device('cuda:1') 如果我们有多个GPU的话，可以在":"后面指定访问哪个GPU.或者使用torch.cuda.device(i)返回第i个gpu
torch.device('cpu'), torch.device('cuda')

(device(type='cpu'), device(type='cuda'))

# 查看GPU的数量
torch.cuda.device_count()

1

# 这两个函数允许我们在请求GPU不存在的情况下运行代码
def try_gpu(i=0):
    """如果存在GPU就返回GPU(i)，否则返回cpu()"""
    if torch.cuda.device_count() >= i+1:
        return torch.device(i)
    return torch.device('cpu')
def try_all_gpus():
    """返回所有可用的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()

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

# 查询张量所在的设，默认是在cpu上
x = torch.tensor([1, 2, 3])
x.device

device(type='cpu')

# 存储在gpu上
X = torch.ones(2, 3, device=try_gpu())  # 使用device来指定放置在gpu上
X

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

# 在第0张GPU上存值
Y = torch.rand(2, 3, device=try_gpu(0))
Y

tensor([[0.7585, 0.3336, 0.2373],
        [0.6617, 0.2012, 0.7119]], device='cuda:0')

# 要计算X + Y，我们需要决定在哪里执行这个操作，我们必须保证X、Y在同一个gpu上
Z = X.cuda(0)
print(X)
print(Z)

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

# 在同一个GPU上的数据我们可以计算
Y + Z

tensor([[1.7585, 1.3336, 1.2373],
        [1.6617, 1.2012, 1.7119]], device='cuda:0')

Z.cuda(0) is Z  # 不会在同一张gpu上拷贝自己

True

# 神经网络与GPU
net = nn.Sequential(nn.Linear(3, 1))  # 创建神经网络
net = net.to(device=try_gpu())  # 使用to函数将创建好的神经网络挪到某个gpu上
net(X)

tensor([[-0.1770],
        [-0.1770]], device='cuda:0', grad_fn=<ThAddmmBackward>)

# 确认模型参数存储在同一额GPU上net[0].weight.data.device

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