GPU
我们先看看如何使用单个NVIDIA GPU进行计算。 首先，确保至少安装了一个NVIDIA GPU。 然后，下载NVIDIA驱动和CUDA 并按照提示设置适当的路径。 当这些准备工作完成，就可以使用nvidia-smi命令来查看显卡信息。
!nvidia-smi
Sun Oct 19 23:28:13 2025
+---------------------------------------------------------------------------------------+
| NVIDIA-SMI 535.230.02 Driver Version: 535.230.02 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 NVIDIA GeForce RTX 3060 Ti Off | 00000000:01:00.0 On | N/A |
| 32% 36C P5 20W / 200W | 927MiB / 8192MiB | 20% Default |
| | | N/A |
+-----------------------------------------+----------------------+----------------------+
+---------------------------------------------------------------------------------------+
| Processes: |
| GPU GI CI PID Type Process name GPU Memory |
| ID ID Usage |
|=======================================================================================|
| 0 N/A N/A 1709 G /usr/lib/xorg/Xorg 453MiB |
| 0 N/A N/A 1938 G /usr/bin/gnome-shell 57MiB |
| 0 N/A N/A 6458 G ...irefox/7084/usr/lib/firefox/firefox 192MiB |
| 0 N/A N/A 30296 G ...AAAAAAAACAAAAAAAAAA= --shared-files 22MiB |
| 0 N/A N/A 31817 G ...erProcess --variations-seed-version 149MiB |
+---------------------------------------------------------------------------------------+
在PyTorch中，每个数组都有一个设备（device）， 我们通常将其称为环境（context）。 默认情况下，所有变量和相关的计算都分配给CPU。 有时环境可能是GPU。 当我们跨多个服务器部署作业时，事情会变得更加棘手。 通过智能地将数组分配给环境， 我们可以最大限度地减少在设备之间传输数据的时间。 例如，当在带有GPU的服务器上训练神经网络时， 我们通常希望模型的参数在GPU上。
在PyTorch中，CPU和GPU可以用torch.device('cpu') 和torch.device('cuda')表示。 应该注意的是，cpu设备意味着所有物理CPU和内存， 这意味着PyTorch的计算将尝试使用所有CPU核心。 然而，gpu设备只代表一个卡和相应的显存。 如果有多个GPU，我们使用torch.device(f'cuda:{i}') 来表示第块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()
现在我们定义了两个方便的函数， 这两个函数允许我们在不存在所需所有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)])
我们可以查询张量所在的设备
x = torch.tensor([1, 2, 3])
x.device
(device(type='cpu'), -1)
需要注意的是，无论何时我们要对多个项进行操作， 它们都必须在同一个设备上。 例如，如果我们对两个张量求和， 我们需要确保两个张量都位于同一个设备上， 否则框架将不知道在哪里存储结果，甚至不知道在哪里执行计算。
存储在GPU上
有几种方法可以在GPU上存储张量。 例如，我们可以在创建张量时指定存储设备。接 下来，我们在第一个gpu上创建张量变量X。 在GPU上创建的张量只消耗这个GPU的显存。 我们可以使用nvidia-smi命令查看显存使用情况。 一般来说，我们需要确保不创建超过GPU显存限制的数据。
X = torch.ones(2, 3, device=try_gpu())
X
tensor([[1., 1., 1.],
[1., 1., 1.]], device='cuda:0')
复制
如果我们要计算X + Y，我们需要决定在哪里执行这个操作。我们需要将这两个数放置在同一个device上。
类似地，神经网络模型可以指定设备。 下面的代码将模型参数放在GPU上。
net = nn.Sequential(nn.Linear(3, 1))
net = net.to(device=try_gpu())
参数
net[0].weight.data.device, net[0].bias.data.device
(device(type='cuda', index=0), device(type='cuda', index=0))