FAFormula and AlgorithmCNNConvolution Netural Network1. 前言在深度学习训练中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.35CUDA Version驱动支持的最高 CUDA 版本例如11.4Memory-Usage显存占用例如2535MiB / 16384MiBGPU-UtilGPU 利用率例如10%。如果命令报错或没有输出 GPU 信息说明驱动或显卡环境存在问题需要先解决环境问题再继续。3. 用 PyTorch 查询设备确认硬件没问题后就可以在 PyTorch 中查询设备了。importtorchfromtorchimportnn torch.device(cpu),torch.cuda.device(cuda),torch.cuda.device(cuda:1)输出(device(typecpu), 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(i0):如果存在则返回 gpu(i)否则返回 cpu()。iftorch.cuda.device_count()i1:returntorch.device(fcuda:{i})returntorch.device(cpu)deftry_all_gpus():返回所有可用的 GPU如果没有 GPU则返回 [cpu()]。devices[torch.device(fcuda:{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(typecuda, index0)]可以看到try_gpu()返回cuda:0try_gpu(10)因为不存在 10 号 GPU自动回退到cputry_all_gpus()返回所有可用 GPU 的列表。这种写法让代码具备良好的可移植性同一份脚本既能在 GPU 服务器上跑也能在普通笔记本上跑。5. 张量的设备管理5.1 查询张量所在设备每个张量都有一个device属性表示它当前存储在哪里Xtorch.tensor([1,2,3])print(X.device)# 默认在 CPU 内存上输出cpu5.2 在 GPU 上创建张量创建张量时通过device参数直接指定设备Xtorch.ones(2,3,devicetry_gpu())X输出tensor([[1., 1., 1.], [1., 1., 1.]], devicecuda:0)5.3 设备不存在时的回退如果指定的 GPU 不存在try_gpu会返回 CPU张量就落在 CPU 上Ytorch.rand(2,3,devicetry_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直接相加会报错。需要先把它们统一到同一设备ZX.cuda(0)# X Y 必须 X 和 Y 都在同一个 GPU 上print(X)print(Z)输出tensor([[1., 1., 1.], [1., 1., 1.]], devicecuda:0) tensor([[1., 1., 1.], [1., 1., 1.]], devicecuda:0)把Y也放到 0 号 GPU 后即可相加Ytorch.rand(2,3,devicetry_gpu(0))YZ另外.cuda(0)对已经在 0 号 GPU 上的张量是幂等的不会重复拷贝Z.cuda(0)isZ# 如果变量在 0 号 GPU 时就返回 True输出True6. 把神经网络放到 GPU 上模型本身也是由张量参数组成的因此同样需要迁移到 GPU。推荐使用.to(device)的写法netnn.Sequential(nn.Linear(3,1))# 创建神经网络时已经把权重初始化好了netnet.to(devicetry_gpu())# 把所有参数在 0 号 GPU 上拷贝一份Xtorch.ones(2,3,devicetry_gpu())# X 在 0 号 GPU 上net(X)# 前向运算的所有元素都在 0 号 GPU 上运行输出tensor([[0.7668], [0.7668]], devicecuda:0, grad_fnAddmmBackward0)确认模型参数确实在 GPU 上net[0].weight.data.device输出device(typecuda, index0)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 之间自由切换。