GPU的使用
因为电脑是集显,所以选择使用colab。
!nvidia-smi
Sun Apr 7 05:48:37 2024
+---------------------------------------------------------------------------------------+
| NVIDIA-SMI 535.104.05 Driver Version: 535.104.05 CUDA Version: 12.2 |
|-----------------------------------------+----------------------+----------------------+
| GPU Name Persistence-M | Bus-Id Disp.A | Volatile Uncorr. ECC |
| Fan Temp Perf Pwr:Usage/Cap | Memory-Usage | GPU-Util Compute M. |
| | | MIG M. |
|=========================================+======================+======================|
| 0 Tesla T4 Off | 00000000:00:04.0 Off | 0 |
| N/A 44C P8 10W / 70W | 0MiB / 15360MiB | 0% Default |
| | | N/A |
+-----------------------------------------+----------------------+----------------------+
+---------------------------------------------------------------------------------------+
| Processes: |
| GPU GI CI PID Type Process name GPU Memory |
| ID ID Usage |
|=======================================================================================|
| No running processes found |
+---------------------------------------------------------------------------------------+
cpu设备意味着所有物理CPU和内存, 这意味着PyTorch的计算将尝试使用所有CPU核心。
然而,gpu设备只代表一个卡和相应的显存。
如果有多个GPU,我们使用torch.device(f’cuda:{i}') 来表示第0块GPU(从0开始)。
另外,cuda:0和cuda是等价的。
import torch
from torch import nn
torch.device('cpu'), torch.device('cuda'), torch.device('cuda:1')
(device(type='cpu'), device(type='cuda'), device(type='cuda', index=1))
查看可用gpu的数量
torch.cuda.device_count()
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()
(device(type='cuda', index=0),
device(type='cpu'),
[device(type='cuda', index=0)])
在gpu上创建tensor
查看张量所在的设备
x=torch.tensor([1,2,3])
x.device
device(type='cpu')
存储在gpu上
X=torch.ones(2,3,device=try_gpu()) # device:返回第0个gpu,即张量存储在第0个gpu上
X
tensor([[1., 1., 1.],
[1., 1., 1.]], device='cuda:0')
第二个gpu上创建一个随机张量
Y=torch.rand(2,3,device=try_gpu(1)) # device:返回第1个gpu,即张量存储在第1个gpu上
Y
tensor([[0.6012, 0.2272, 0.6751],
[0.5924, 0.3950, 0.2923]])
要计算X+Y,需要决定在哪里执行这个操作
Z = Y.cuda() # 将Y移动到第0个gpu,将X移动到第1个会报错
print(Y)
print(Z)
tensor([[0.6012, 0.2272, 0.6751],
[0.5924, 0.3950, 0.2923]])
tensor([[0.6012, 0.2272, 0.6751],
[0.5924, 0.3950, 0.2923]], device='cuda:0')
神经网络与GPU
net=nn.Sequential(nn.Linear(3,1))
net=net.to(device=try_gpu()) #将所有的参数在第0个GPU上copy一份
net(X) # X也是在第一个GPU,因此前向运算也在第0个GPU
tensor([[-0.3243],
[-0.3243]], device='cuda:0', grad_fn=<AddmmBackward0>)
确认模型参数存储在同一个GPU上
net[0].weight.data.device
device(type='cuda', index=0)
只要所有的数据和参数都在同一个设备上,我们就可以有效地学习模型。