FA：Formula and Algorithm，CNN：Convolution Netural Network
1. 前言
在深度学习训练中，GPU 往往是决定训练速度的关键。PyTorch 提供了一套简洁而灵活的 GPU 使用方式：通过torch.device抽象设备，用.to()把张量和模型搬到指定设备上。
本文从最基础的 GPU 确认开始，逐步介绍设备查询、张量迁移、多 GPU 兼容写法，以及如何把神经网络放到 GPU 上运行，帮助你写出既高效又健壮的 PyTorch 代码。
2. 确认 GPU 是否可用
第一步永远是确认环境里有没有 GPU、驱动和 CUDA 版本是否正常。在 Jupyter Notebook 中可以直接执行：
!nvidia-smi输出类似下面这样：
Sat Apr 30 18:03:10 2022 +-----------------------------------------------------------------------------+ | NVIDIA-SMI 471.35 Driver Version: 471.35 CUDA Version: 11.4 | |-------------------------------+----------------------+----------------------+ | GPU Name TCC/WDDM | Bus-Id Disp.A | Volatile Uncorr. ECC | | Fan Temp Perf Pwr:Usage/Cap| Memory-Usage | GPU-Util Compute M. | |===============================+======================+======================| | 0 NVIDIA GeForce ... WDDM | 00000000:01:00.0 On | N/A | | N/A 57C P5 27W / N/A | 2535MiB / 16384MiB | 10% Default | +-------------------------------+----------------------+----------------------+从这张表里可以读出几个关键信息：
- Driver Version：显卡驱动版本，例如
471.35； - CUDA Version：驱动支持的最高 CUDA 版本，例如
11.4； - Memory-Usage：显存占用，例如
2535MiB / 16384MiB； - GPU-Util：GPU 利用率，例如
10%。
如果命令报错或没有输出 GPU 信息，说明驱动或显卡环境存在问题，需要先解决环境问题再继续。
3. 用 PyTorch 查询设备
确认硬件没问题后，就可以在 PyTorch 中查询设备了。
importtorchfromtorchimportnn torch.device('cpu'),torch.cuda.device('cuda'),torch.cuda.device('cuda:1')输出：
(device(type='cpu'), <torch.cuda.device at 0x2a59c8b2278>, <torch.cuda.device at 0x2a59c8b2588>)torch.device只是一个设备描述对象，即使机器上不存在 1 号 GPU，也可以创建cuda:1这样的描述，真正使用时会报错。
查看当前可用 GPU 数量：
print(torch.cuda.device_count())输出：
1说明当前环境只有一块可用 GPU。
4. 封装设备选择函数
为了让代码在「有 GPU」和「无 GPU」两种环境下都能运行，推荐封装两个工具函数：
deftry_gpu(i=0):"""如果存在，则返回 gpu(i)，否则返回 cpu()。"""iftorch.cuda.device_count()>=i+1:returntorch.device(f'cuda:{i}')returntorch.device('cpu')deftry_all_gpus():"""返回所有可用的 GPU，如果没有 GPU，则返回 [cpu()]。"""devices=[torch.device(f'cuda:{i}')foriinrange(torch.cuda.device_count())]returndevicesifdeviceselse[torch.device('cpu')]测试一下：
print(try_gpu())print(try_gpu(10))print(try_all_gpus())输出：
cuda:0 cpu [device(type='cuda', index=0)]可以看到：
try_gpu()返回cuda:0；try_gpu(10)因为不存在 10 号 GPU，自动回退到cpu；try_all_gpus()返回所有可用 GPU 的列表。
这种写法让代码具备良好的可移植性，同一份脚本既能在 GPU 服务器上跑，也能在普通笔记本上跑。
5. 张量的设备管理
5.1 查询张量所在设备
每个张量都有一个device属性，表示它当前存储在哪里：
X=torch.tensor([1,2,3])print(X.device)# 默认在 CPU 内存上输出：
cpu5.2 在 GPU 上创建张量
创建张量时通过device参数直接指定设备：
X=torch.ones(2,3,device=try_gpu())X输出：
tensor([[1., 1., 1.], [1., 1., 1.]], device='cuda:0')5.3 设备不存在时的回退
如果指定的 GPU 不存在，try_gpu会返回 CPU，张量就落在 CPU 上：
Y=torch.rand(2,3,device=try_gpu(1))print(Y.device)# 没有 1 号 GPU，则放到 CPU 上Y输出：
cpu tensor([[0.3045, 0.4398, 0.9162], [0.5351, 0.1256, 0.8916]])5.4 跨设备运算的限制
PyTorch 要求参与运算的张量位于同一设备上。如果X在 GPU、Y在 CPU，直接相加会报错。需要先把它们统一到同一设备：
Z=X.cuda(0)# X + Y 必须 X 和 Y 都在同一个 GPU 上print(X)print(Z)输出：
tensor([[1., 1., 1.], [1., 1., 1.]], device='cuda:0') tensor([[1., 1., 1.], [1., 1., 1.]], device='cuda:0')把Y也放到 0 号 GPU 后即可相加：
Y=torch.rand(2,3,device=try_gpu(0))Y+Z另外，.cuda(0)对已经在 0 号 GPU 上的张量是幂等的，不会重复拷贝：
Z.cuda(0)isZ# 如果变量在 0 号 GPU 时，就返回 True输出：
True6. 把神经网络放到 GPU 上
模型本身也是由张量（参数）组成的，因此同样需要迁移到 GPU。推荐使用.to(device)的写法：
net=nn.Sequential(nn.Linear(3,1))# 创建神经网络时已经把权重初始化好了net=net.to(device=try_gpu())# 把所有参数在 0 号 GPU 上拷贝一份X=torch.ones(2,3,device=try_gpu())# X 在 0 号 GPU 上net(X)# 前向运算的所有元素都在 0 号 GPU 上运行输出：
tensor([[0.7668], [0.7668]], device='cuda:0', grad_fn=<AddmmBackward0>)确认模型参数确实在 GPU 上：
net[0].weight.data.device输出：
device(type='cuda', index=0)7. 实践建议
- 统一使用
.to(device)：相比.cuda()，.to(device)更通用，切换 CPU/GPU 时只需改一处。 - 输入与模型同设备：前向计算前务必确认输入张量和模型参数在同一设备上，否则会抛出设备不匹配错误。
- 善用
try_gpu封装：避免在代码里硬编码cuda:0，让脚本在无 GPU 环境下也能跑通。 - 注意显存占用：
nvidia-smi中的Memory-Usage能帮助你判断是否出现显存泄漏或批次过大。 - 多卡场景：
try_all_gpus()返回设备列表，可配合DataParallel或DistributedDataParallel使用。
8. 小结
本文从nvidia-smi确认 GPU 开始，介绍了 PyTorch 中设备查询、张量迁移、跨设备运算限制，以及模型上 GPU 的完整流程。核心思路只有一条：用torch.device描述设备，用.to(device)统一迁移，用try_gpu保证代码可移植。掌握这套写法后，你的 PyTorch 代码就能在 CPU 与 GPU 之间自由切换。
 】Docker与容器（上篇）)
