PyTorch训练SegNet时数据未入GPU引发RuntimeError问题排查

问题根源 #

错误提示输入张量是CPU的torch.FloatTensor，但模型权重是GPU的torch.cuda.FloatTensor，说明计算训练/验证分数时，输入数据没有转移到GPU。

具体来看：

• 训练阶段你已经将X_batch和Y_batch转到device，所以训练过程正常完成并输出train_loss。
• 但执行avg_score_train = score_fn(model, iou_pytorch, data_tr)时，data_tr返回的张量仍在CPU上，而模型参数在GPU上，此时模型接收CPU输入就会触发类型不匹配错误。
• Colab上可能因为默认数据加载时自动适配设备，或者你的score_fn在Colab环境中隐式处理了设备转移，所以未报错。

修复步骤 #

1. 修正score_fn函数 #

确保score_fn内部将输入数据转移到指定设备。修改后的score_fn示例：

def score_fn(model, metric_fn, dataloader, device):
    model.eval()
    total_score = 0.0
    with torch.no_grad():
        for X, Y in dataloader:
            # 将数据转移到指定设备
            X = X.to(device)
            Y = Y.to(device)
            Y_pred = model(X)
            score = metric_fn(Y_pred, Y)
            total_score += score / len(dataloader)
    return total_score

2. 调整训练函数中的score_fn调用 #

在训练函数中调用score_fn时传入device参数：

# train score
avg_score_train = score_fn(model, iou_pytorch, data_tr, device)
scores_train.append(avg_score_train)

# ...

# val score
avg_score_val = score_fn(model, iou_pytorch, data_vl, device)
scores_val.append(avg_score_val)

3. 优化验证阶段的张量转移（可选但推荐） #

验证阶段当前将Y_hat转回CPU后和CPU上的Y_val计算损失，建议统一在GPU上完成计算，避免不必要的张量转移，提升效率：

# val loss
avg_loss_val = 0
model.eval()  # testing mode
for X_val, Y_val in data_vl:
    with torch.no_grad():
        X_val = X_val.to(device)
        Y_val = Y_val.to(device)            
        Y_hat = model(X_val)
        loss = loss_fn(Y_hat, Y_val) # 直接在GPU上计算损失
        avg_loss_val += loss / len(data_vl)

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