116_深度学习提速：PyTorch 利用 GPU 训练的全攻略

在 PyTorch 中利用 GPU 训练并不复杂核心逻辑就是将模型和数据“搬”到显卡内存中。你的笔记中详细记录了两种实现方式下面我们结合代码逐一拆解。1. 哪些部分需要移动到 GPU要实现 GPU 加速必须保证以下三部分都在同一块显卡上网络模型(nn.Module)损失函数(Loss Function)数据(包括输入imgs和标注targets)2. 方式一使用.cuda()(基础版)这种方式简单直接通过调用对象的.cuda()方法来实现迁移。代码实战Python# 1. 网络模型转移到 GPU tudui Tudui() if torch.cuda.is_available(): tudui tudui.cuda() # 2. 损失函数转移到 GPU loss_fn nn.CrossEntropyLoss() if torch.cuda.is_available(): loss_fn loss_fn.cuda() # 3. 训练循环中的数据转移 for data in train_dataloader: imgs, targets data if torch.cuda.is_available(): imgs imgs.cuda() targets targets.cuda() outputs tudui(imgs) loss loss_fn(outputs, targets) # ... 后续优化逻辑注意使用.cuda()前务必先用torch.cuda.is_available()判断当前环境是否有显卡。3. 方式二使用.to(device)(推荐方案)这种方式更加通用且优雅。你可以定义一个device变量后续一键切换 CPU 或 GPU。代码实战Python# 定义训练设备 device torch.device(cuda if torch.cuda.is_available() else cpu) # 1. 转移模型 tudui Tudui() tudui tudui.to(device) # 2. 转移损失函数 loss_fn nn.CrossEntropyLoss() loss_fn loss_fn.to(device) # 3. 转移数据 for data in train_dataloader: imgs, targets data imgs imgs.to(device) targets targets.to(device) # ... 执行训练4. 关键点深度解析为什么模型和损失函数不需要重新赋值而数据需要在你的代码中你可能会发现tudui.to(device)执行后模型就变了。但数据必须写成imgs imgs.to(device)。原因模型nn.Module的.to()方法会直接修改其内部参数而张量Tensor的.to()方法会返回一个新的副本所以必须重新赋值。如何查看显卡占用情况文件中最后提到通过命令行查看 GPU 状态Bash!nvidia-smi通过这个命令你可以实时监控显存占用Memory-Usage和 GPU 利用率Volatile GPU-Util确保你的显卡正在全力工作。5. 总结GPU 训练的“避坑”准则保持一致性输入数据和模型必须在同一个 device 上否则会报RuntimeError: Expected all tensors to be on the same device。单机多卡处理如果你有多个 GPU可以通过cuda:0,cuda:1来指定特定的显卡。内存管理如果遇到out of memory(OOM)尝试调小batch_size。 学习小结学会利用 GPU 训练后你已经解锁了处理大规模图像任务的能力。