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

在 PyTorch 中，利用 GPU 训练并不复杂，核心逻辑就是：将模型和数据"搬"到显卡内存中。 你的笔记中详细记录了两种实现方式，下面我们结合代码逐一拆解。

1. 哪些部分需要移动到 GPU？

要实现 GPU 加速，必须保证以下三部分都在同一块显卡上：

1. 网络模型 (nn.Module)
2. 损失函数 (Loss Function)
3. 数据 (包括输入 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 训练的"避坑"准则

1. 保持一致性 ：输入数据和模型必须在同一个 device 上，否则会报 RuntimeError: Expected all tensors to be on the same device。
2. 单机多卡处理 ：如果你有多个 GPU，可以通过 cuda:0, cuda:1 来指定特定的显卡。
3. 内存管理 ：如果遇到 out of memory (OOM)，尝试调小 batch_size。

---

💡 学习小结

学会利用 GPU 训练后，你已经解锁了处理大规模图像任务的能力。