PyTorch线性回归验证损失为NaN问题求助

问题根源 #

警告明确指出目标张量尺寸([128,5])与模型输出张量尺寸([128,1])不匹配，这会触发广播机制导致loss计算异常，最终出现val_loss: nan。大概率是数据处理时误将目标列（保费charges）当成分类列做了编码，或者提取输出时取错了列。

具体修复步骤 #

1. 修正数据处理流程 #

确保仅对输入特征中的分类列编码，目标列单独提取且保持一维：

import pandas as pd
import numpy as np
import torch
from torch.utils.data import TensorDataset, DataLoader, random_split

# 加载数据
df = pd.read_csv("https://gist.github.com/BirajCoder/5f068dfe759c1ea6bdfce9535acdb72d/raw/c84d84e3c80f93be67f6c069cbdc0195ec36acbd/insurance.csv")
df_copy = df.copy()

# 仅对输入分类列编码（sex, smoker, region）
df_encoded = pd.get_dummies(df_copy, columns=['sex', 'smoker', 'region'], drop_first=True)

# 提取输入特征（除charges列）和目标值（仅charges列）
X = df_encoded.drop('charges', axis=1).values
y = df_encoded['charges'].values.reshape(-1, 1)  # 转为[N,1]形状，匹配模型输出

# 转为PyTorch张量
X_tensor = torch.tensor(X, dtype=torch.float32)
y_tensor = torch.tensor(y, dtype=torch.float32)

2. 确认数据集与DataLoader的正确性 #

创建TensorDataset时确保输入和目标的维度严格匹配：

# 创建数据集
dataset = TensorDataset(X_tensor, y_tensor)

# 拆分训练集（80%）和验证集（20%）
train_size = int(0.8 * len(dataset))
val_size = len(dataset) - train_size
train_dataset, val_dataset = random_split(dataset, [train_size, val_size])

# 创建DataLoader
batch_size = 128
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False)

3. 修正模型输出维度 #

确保InsuranceModel的最后一层输出为1维（对应回归任务的单个预测值）：

import torch.nn as nn
import torch.nn.functional as F

class InsuranceModel(nn.Module):
    def __init__(self, input_features):
        super().__init__()
        self.linear = nn.Linear(input_features, 1)  # 输出维度固定为1
    
    def forward(self, x):
        return self.linear(x)

# 初始化模型：输入特征数为编码后的特征列总数
input_features = df_encoded.drop('charges', axis=1).shape[1]
model = InsuranceModel(input_features)

4. 确保训练/验证逻辑中loss计算的维度匹配 #

在fit函数中，保证模型输出与目标张量的形状完全一致：

def fit(epochs, lr, model, train_loader, val_loader):
    optimizer = torch.optim.SGD(model.parameters(), lr=lr)
    history = {'train_loss': [], 'val_loss': []}
    
    for epoch in range(epochs):
        # 训练阶段
        model.train()
        train_loss = 0.0
        for batch in train_loader:
            inputs, targets = batch
            optimizer.zero_grad()
            outputs = model(inputs)
            loss = F.mse_loss(outputs, targets)  # 此时outputs和targets均为[batch_size,1]
            loss.backward()
            optimizer.step()
            train_loss += loss.item() * inputs.size(0)
        
        train_loss /= len(train_loader.dataset)
        history['train_loss'].append(train_loss)
        
        # 验证阶段
        model.eval()
        val_loss = 0.0
        with torch.no_grad():
            for batch in val_loader:
                inputs, targets = batch
                outputs = model(inputs)
                loss = F.mse_loss(outputs, targets)
                val_loss += loss.item() * inputs.size(0)
        
        val_loss /= len(val_loader.dataset)
        history['val_loss'].append(val_loss)
        print(f"Epoch {epoch+1}/{epochs}, train_loss: {train_loss:.4f}, val_loss: {val_loss:.4f}")
    
    return history

5. 可选：调整学习率避免数值溢出 #

如果修复维度后仍出现nan，可以尝试降低学习率（比如从1e-2调到1e-3），较大的学习率可能导致权重更新过大，引发数值溢出。

验证效果 #

重新运行训练代码：

epochs = 30
lr = 1e-3  # 可选调整
history1 = fit(epochs, lr, model, train_loader, val_loader)

此时维度匹配警告会消失，val_loss也会正常输出数值而非nan。

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