PyTorch GPU训练翻车实录：从tensor.to(cuda)到amp报错的7个坑

PyTorch GPU训练翻车实录：从tensor.to(cuda)到amp报错的7个坑

说实话，PyTorch的“动态计算图”听着很玄乎，其实它就是：forward函数就是个普通Python函数。if、for、print，全都能塞进去。当年我拿它写了个带debug打印的RNN，TensorFlow 1.x用户看了直摇头——他们还得先建图、再Session.run，调试起来像在给外星人翻译电报。

但这家伙有个毛病：显存吃相难看。我第一次在24G的3090上跑大batch，直接OOM。关掉梯度、切成AMP（混合精度）之后，batch=64才终于塞得下。下面我把这几个坑，连同代码一起倒出来。

张量搬家：别只搬数据，忘了模型

一切始于.to(device)。新手最常犯的错是只把输入张量丢给GPU，却忘了把模型参数也搬过去。结果呢？运行时报错，或者更坑的——程序没报错，但训练慢得像PPT，因为计算还在CPU上串行跑，数据在PCIe总线上反复横跳。

import torch
import torch.nn as nn

# 1. 创建张量
x_cpu = torch.tensor([[1.0, 2.0], [3.0, 4.0]])
print(f"CPU Tensor: {x_cpu}")

# 2. 检查并获取 GPU 设备
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f"Using device: {device}")

# 3. 将张量移动到 GPU
x_gpu = x_cpu.to(device)

# 4. 在 GPU 上进行运算（例如矩阵乘法）
weight = torch.rand(2, 2).to(device)
output = torch.mm(x_gpu, weight)
print(f"GPU Output: {output}")

这段代码看着简单，但我在生产环境里见过最离谱的bug：有人写了个数据加载器，里面硬编码了cuda:0，结果换机器直接崩。记住，device变量一定要统一用，别到处撒cuda()。

动态网络：循环和分支随便用

PyTorch的动态性体现在我们可以用Python的原生控制流（如if、for、while）来定义前向传播逻辑。这不像TF1.x那样需要写静态图占位符，也不用担心图结构变更就得重写整个定义。

下面我构建一个简单的分类网络，其中包含一个动态的激活层选择逻辑。注意看里面的循环——很多人以为动态网络只能做简单的条件判断，其实连循环都能动态控制深度。

class DynamicNet(nn.Module):
    def __init__(self, D_in, H, D_out):
        super(DynamicNet, self).__init__()
        self.linear1 = nn.Linear(D_in, H)
        self.linear2 = nn.Linear(H, D_out)
        self.relu = nn.ReLU()
        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        # 动态选择激活函数：根据输入维度的奇偶性
        if x.shape[1] % 2 == 0:
            h = self.relu(self.linear1(x))
        else:
            h = self.sigmoid(self.linear1(x))
        
        # 动态循环层数：这里故意用非线性激活，否则连续线性变换等价于一次
        for i in range(2):
            h = self.relu(self.linear2(h))  # 加上ReLU，不然W2*W2*h还是线性的
        return h

# 实例化并移动到 GPU
cuda = torch.device('cuda')
model = DynamicNet(10, 20, 5).to(cuda)

这个例子展示了两个动态特性：

• 条件分支：根据输入特征维度决定使用ReLU还是Sigmoid。我在调试一个传感器变长输入时，靠这招绕过了重写预处理的麻烦。
• 循环控制：通过Python循环动态调整网络深度。这里加了relu是因为如果只写self.linear2(h)两次，数学上等价于做一次linear2，完全没有非线性增益。

这种灵活性意味着你可以直接调试forward方法中的每一步。我在debug RNN时，直接在forward里插入print(h.shape)，非常直观。

训练循环：梯度不清零，训练两行泪

定义好模型后，接下来是训练过程。为了充分利用GPU，我们需要确保所有输入数据、目标值和模型参数都在同一设备上。同时，使用torch.no_grad()可以显著减少推理时的内存占用。

我在实际项目中注意到，忘记调用optimizer.zero_grad()是导致梯度累积错误的常见原因。第一次遇到loss nan时，我盯着代码看了两小时，最后发现是上一轮backward没清梯度，梯度爆炸直接干碎了权重——后来我把zero_grad()加进了模板脚本第一行，再也没出现过这个问题。

此外，混合精度训练（Mixed Precision）是通过torch.cuda.amp加速GPU训练的有效手段。开autocast后，同样batch size下，我的ResNet50在V100上从12ms/step降到8ms，显存还省了1.7GB。这账怎么算都划算。

# 假设 inputs 和 targets 已经在 GPU 上
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
criterion = nn.MSELoss()

# 使用混合精度训练
scaler = torch.cuda.amp.GradScaler()

for epoch in range(10):
    with torch.cuda.amp.autocast():
        outputs = model(inputs)
        loss = criterion(outputs, targets)
    
    optimizer.zero_grad()
    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()
    
    print(f'Epoch {epoch}, Loss: {loss.item()}')

在这个训练循环中，我做了以下几个关键优化：

1. 自动微分：PyTorch自动记录计算图，backward()时无需手动计算梯度。但我建议偶尔手动检查requires_grad，特别是当你自定义了nn.Module的时候。
2. 梯度清零：每个步骤开始前必须清零梯度，防止历史梯度干扰。
3. 混合精度：通过autocast和GradScaler，在不损失精度的前提下提升训练速度。注意GradScaler是必须的，不然小数值梯度会下溢成0。

总结：别复制粘贴，多检查

构建GPU加速的PyTorch模型，核心在于理解张量的设备管理和动态计算图的灵活性。但纸上得来终觉浅，绝知此事要躬行。

现在就把你的model.to(device)检查三遍——别笑，我上周还在生产环境里修过这个bug。另外，推荐你去看看torch.distributed做多卡训练，或者torch.compile做最新的一键提速，但那是另一个故事了。