PyTorch模型GPU训练性能远逊于CPU的问题求助

在CPU上训练的小型PyTorch模型验证准确率约0.91，但训练速度较慢。仅对代码做GPU适配修改后在Google Colab的GPU上运行（其余代码完全一致），尽管训练轮次更多，模型验证准确率仅约0.71，性能大幅下降。尝试Colab提供的TPU时也存在同样问题，运行环境为Google Drive，怀疑是浮点精度问题。

代码示例 #

CPU训练代码 #

# Train the false colour model
epochs = 5
loss_func = nn.MSELoss()
optimizer = optim.AdamW(params = mymodel_false.parameters(), lr = 0.0004)

losses = []
for epoch in range(epochs):
    for i, (input_img, target_img) in enumerate(train_loader_false):
        # Zero the gradients
        optimizer.zero_grad()

        # forward + backward + optimize
        output_img = mymodel_false(input_img)
        loss = loss_func(output_img, target_img)
        loss.backward()
        optimizer.step()

        losses.append(loss)
        print(f"Epoch: {epoch}, Loss: {loss}")

GPU训练代码 #

# Train the false colour model
device = torch.device("cuda")
model = model.to(device)
epochs = 5
loss_func = nn.MSELoss()
optimizer = optim.AdamW(params = mymodel_false.parameters(), lr = 0.0004)

losses = []
for epoch in range(epochs):
    for i, (input_img, target_img) in enumerate(train_loader_false):
        # Zero the gradients
        optimizer.zero_grad()

        # forward + backward + optimize
        input_img = input_img.to(device)
        output_img = mymodel_false(input_img)
        target_img = target_img.to(device)
        loss = loss_func(output_img, target_img)
        loss.backward()
        optimizer.step()

        losses.append(loss)
        print(f"Epoch: {epoch}, Loss: {loss}")

排查与解决方案 #

1. 模型设备迁移的变量名错误（最可能原因） #

GPU代码中，你执行了model = model.to(device)，但后续训练调用的是mymodel_false，这意味着实际训练的还是CPU上的原始模型，GPU上的模型从未被使用。这会导致GPU训练时模型参数始终处于初始状态，自然性能暴跌。

修正方式：将模型迁移代码改为：

mymodel_false = mymodel_false.to(device)

2. 浮点精度一致性检查 #

CPU与GPU的浮点精度默认可能存在差异，若模型、数据的dtype不统一，会导致训练过程中的精度损失：

• 确保模型和所有张量使用相同的dtype（例如torch.float32），迁移设备时显式指定：

  mymodel_false = mymodel_false.to(device, dtype=torch.float32)
  input_img = input_img.to(device, dtype=torch.float32)
  target_img = target_img.to(device, dtype=torch.float32)
  

• 若Colab默认开启了混合精度训练，可能会导致精度丢失，可显式关闭或使用torch.cuda.amp规范处理混合精度：

  from torch.cuda.amp import GradScaler, autocast
  
  scaler = GradScaler()
  # 训练循环内修改：
  optimizer.zero_grad()
  with autocast():
      output_img = mymodel_false(input_img)
      loss = loss_func(output_img, target_img)
  scaler.scale(loss).backward()
  scaler.step(optimizer)
  scaler.update()
  

3. 固定随机种子消除训练随机性差异 #

GPU/TPU训练时，数据加载、底层运算的随机性可能与CPU不同，导致训练数据分布或权重更新轨迹不一致：

import torch
import numpy as np
import random

def set_seed(seed=42):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False

set_seed()

在训练前调用该函数，确保CPU、GPU、TPU训练的随机性一致。

4. TPU适配的额外要求 #

TPU对PyTorch的适配需要依赖torch_xla库，单纯迁移设备无法正常工作，需调整代码：

• 安装torch_xla：

  !pip install torch_xla[tpu] -f https://storage.googleapis.com/tpu-pytorch/wheels/torch_xla-latest-cp310-cp310-linux_x86_64.whl
  

• 调整训练循环适配TPU：

  import torch_xla.core.xla_model as xm
  
  device = xm.xla_device()
  mymodel_false = mymodel_false.to(device)
  
  # 训练循环内优化器步骤需改为：
  loss.backward()
  xm.optimizer_step(optimizer)
  

内容的提问来源于stack exchange，提问作者Calvin Adams