☰
PyTorch DataLoader实战:加载糖尿病数据集训练回归模型完整指南
2026/10/14 9:51:42 网站建设 项目流程

刚开始接触 PyTorch 的时候,我一度觉得 DataLoader 就是个“装数据的大筐”,把样本塞进去、训练时再一批一批取出来,仅此而已。直到后来有一次,我自己处理一个带时序关系的真实业务数据集,因为偷懒没做 shuffle 踩了大坑,才发现这个工具远没有表面看起来那么简单。这次就以一个最经典的入门场景——用 DataLoader 加载糖尿病数据集并训练回归模型——把数据加载、预处理、训练验证这一整条链路,完整走一遍,顺便把踩过的坑和排查经验也一并交代清楚。

这个场景适合三类人:刚入门 PyTorch、想搞清楚 Dataset 和 DataLoader 到底怎么用的新手;已经能跑通分类任务、但没碰过回归任务的进阶者;以及想把数据处理流程规范化、避免“代码能跑但结果稀烂”的老手。如果你正在学机器学习,这个数据集和这套流程,就是你从“调包侠”走向“懂原理”的一座桥。

1. 这个任务到底在做什么:数据、场景与价值的全貌

1.1 糖尿病数据集究竟长什么样

标题里提到的“糖尿病数据集”,绝大多数情况下指的是 scikit-learn 内置的load_diabetes()数据集。它一共有 442 个样本,每个样本包含 10 个特征,分别代表年龄、性别、体质指数(BMI)、平均血压,以及六项血清指标。目标值是一个连续变量,表示在基线状态之后一年,病情进展的定量测量结果。

注意,这不是一个分类任务——不是判断一个人“有没有糖尿病”,而是预测病情进展的数值。换句话说,这是一个典型的回归任务。很多初学者看到“糖尿病”三个字就默认是二分类,结果一打开数据发现目标值是浮点数,当场就懵了。这里必须强调:这个数据集是回归任务,损失函数要用均方误差(MSE),评估指标要看 R² 或 RMSE,而不是准确率。

还有一个容易被忽略的细节:这个数据集的 10 个特征已经做过中心化和缩放处理。什么意思?就是每个特征都已经减去均值、除以标准差(或做了类似的标准归一化),特征之间量纲基本一致。这看起来很方便,但实际使用中反而容易让人麻痹大意,忽略对测试集做独立预处理的问题。后面我会详细讲这一点。

442 个样本,10 个特征,数据量很小。训练一轮模型的时间按秒计算,特别适合用来验证代码链路是否通畅。我个人的经验是:任何新项目、新环境,我都会先拿一份小数据集把训练、验证、评估的完整流程跑通,再上大数据集。糖尿病数据集就是干这个的理想材料。

1.2 为什么拿它当训练样例

很多教程喜欢用 MNIST、CIFAR-10 这类图像数据集做演示,但对于想搞懂回归任务、想理解“特征矩阵 + 连续标签”这种数据形态的人来说,糖尿病数据集反而更合适。原因有三个:

第一,数据形态典型。10 个数值型特征组成特征矩阵,1 个连续目标值组成标签向量。这种形态是表格类机器学习任务的原型,不管是金融风控、医疗预测还是工业参数估计,底色都是这样的数据。把这种数据处理好,换到其他表格数据上可以直接复用。

第二,训练成本极低。442 个样本,两层全连接网络,CPU 上训练几百轮也就几十秒。新手不需要为 GPU 资源发愁,可以反复实验,把注意力集中在数据加载、模型结构、超参数这些真正需要理解的东西上。

第三,回归任务的直观性。预测一个连续值,你可以直接看到预测值和真实值的差距,天然适合理解“模型在拟合什么”。不像分类任务,准确率 80% 你还得去分析混淆矩阵才知道哪儿错了。

顺便说一句,其实在 Kaggle 上还有一个更流行的“糖尿病预测”数据集,那个是二分类任务,包含血糖、血压、皮肤厚度等字段,用来预测是否患病。很多人搜“糖尿病数据集”会搜到那一份,然后和 sklearn 内置这份搞混。我用的是 sklearn 内置的这份回归数据集,下文所有代码都基于它,读者在复现时注意区分。

1.3 数据管道的整体设计思路

整个训练流程,我通常分成四个环节:原始数据获取与划分、预处理与封装、DataLoader 配置、模型训练与评估。每个环节各司其职,边界清晰。

原始数据获取与划分:把数据切成训练集和测试集,一般用 80/20 或 70/30 的比例。划分时要设置随机种子,保证结果可复现。

预处理与封装:这一步最容易出问题。标准做法是在训练集上拟合标准化参数(均值和标准差),然后用同一组参数转换测试集。如果先对整个数据集做标准化再划分,就会造成信息泄露——测试集的信息已经悄悄进入了训练阶段的分布估计,得到的评估指标会偏乐观。

封装与 DataLoader 配置:把训练集和测试集分别包成 Dataset 对象,再传给 DataLoader,设置 batch size、是否打乱、用几个 worker 进程等参数。

模型训练与评估:定义网络结构、损失函数、优化器,跑训练循环,定期在测试集上评估。

这个设计思路本身没什么新意,但绝大部分训练事故,都出在第二步的“信息泄露”和第四步的“评估不当”上。后面的章节,我会一步一步把每个环节的操作细节和原理讲透。

2. DataLoader 机制逐层拆解:不要只会调用,要理解为什么

2.1 Dataset 是原料仓库,DataLoader 是流水线

初学者最容易把 Dataset 和 DataLoader 混为一谈。其实两者的分工非常明确:Dataset 负责“有什么数据”,DataLoader 负责“怎么把数据喂给模型”。

Dataset 是一个抽象类,你需要继承它,并实现两个方法:__len__(返回样本总数)和__getitem__(根据索引返回一个样本)。它就好比一个仓库,你知道里面一共有多少货(__len__),也能根据货架号取出对应的货物(__getitem__)。

DataLoader 则是在仓库外面架设的自动分拣流水线。它按照你指定的规则,从 Dataset 里一批一批地把数据取出来,还能做打乱、并行加载、自动整理成张量等操作。你不需要手动写“用一个 for 循环从第 i 个取到第 i+batch_size 个”这种代码,DataLoader 全帮你干了。

这里有个设计上的好处:Dataset 只关心数据本身,DataLoader 只关心取数据的方式,两者解耦。你可以很方便地替换 Dataset 的实现(比如从文件读取改成从数据库读取),而不影响 DataLoader 的用法;也可以改 DataLoader 的参数(比如 batch size 从 32 改成 64)而不改动数据读取逻辑。

2.2 batch_size、shuffle、num_workers 各管什么

DataLoader 最核心的参数就三个:batch_size、shuffle、num_workers,我逐个拆开讲。

batch_size:每次取多少个样本组成一批。为什么不用全量数据一次性训练?因为小批量训练有两个实际好处:一是内存开销可控,二是每次参数更新只基于一小批数据,自带一定的随机性,有助于跳出局部最优。batch size 太小,梯度估计噪声大,训练不稳定;太大,每个 epoch 更新次数少,收敛变慢且容易陷入尖锐的极小值。对于 442 个样本这种小规模数据集,32 是一个经验上很稳的选择。

shuffle:是否在每个 epoch 开始时打乱数据顺序。训练集必须设为 True,测试集设为 False。原因很简单:如果数据本身存在顺序相关的模式(比如前 100 个样本都是某种类型的病人),不打乱的话模型会学到这种虚假的“顺序规律”,泛化能力变差。我在之前提到的那个业务项目里,就因为没有打乱数据,导致模型在验证集上表现虚高、上线后一塌糊涂。这个教训让我从此把 shuffle 当成铁律。

num_workers:用几个子进程来并行加载数据。数据加载是 I/O 密集型操作,多进程可以显著加速。但注意,不是越大越好。进程切换有开销,数据集太小(比如 442 个样本)时,多进程的收益可以忽略不计,反而可能因为进程通信带来额外负担。糖尿病数据集这种规模,把 num_workers 设为 0 或 2 就够了。

2.3 归一化与数据划分的顺序问题

这个问题值得单独拿出来说,因为几乎每个入门者都会在这里栽跟头。

正确的流程是:先划分训练集和测试集,再在训练集上拟合并应用标准化,最后用训练集的统计量转换测试集。

错误但非常常见的做法是:先对整个数据集做标准化,然后再划分训练集和测试集。这在研究性质的项目里看起来没什么问题,因为标准化参数(均值和标准差)是在全部数据上计算的,测试集的信息已经参与了这些统计量的估计。打个比方:你考试前偷偷看了答案,考试分数自然好看,但真实水平并没有提升。模型评估同理,测试集的作用是模拟“未见过的数据”,一旦它的统计特征渗透进训练过程的任何一个环节,评估结果就失真了。

正确代码长这样:

from sklearn.model_selection import train_test_split from sklearn.preprocessing import StandardScaler X_train, X_test, y_train, y_test = train_test_split( X, y, test_size=0.2, random_state=42 ) scaler = StandardScaler() X_train = scaler.fit_transform(X_train) X_test = scaler.transform(X_test)

注意fit_transform和transform的区别:前者拟合统计量并转换,后者只转换不拟合。测试集只能用训练集拟合好的 scaler 去转换,不能自己再 fit。

2.4 自定义 Dataset 子类的写法要点

虽然借助 sklearn 的train_test_split能拿到划分好的数组,但为了和 DataLoader 正常配合,还需要把这些数组包装成 Dataset。我见过不少新手在这步硬编码,结果代码只能跑一次,换个数据集就全部推倒重来。正确做法是写一个通用的子类:

import torch from torch.utils.data import Dataset class TabularDataset(Dataset): def __init__(self, X, y): self.X = torch.tensor(X, dtype=torch.float32) self.y = torch.tensor(y, dtype=torch.float32).view(-1, 1) def __len__(self): return len(self.X) def __getitem__(self, idx): return self.X[idx], self.y[idx]

有两个细节值得注意。第一,y用view(-1, 1)把形状从(batch_size,)改成(batch_size, 1),这是为了让标签和模型输出形状一致,方便算 MSE 损失。第二,在__init__里就把 numpy 数组转成 torch.Tensor,可以避免在__getitem__里反复做转换开销。442 个样本存不进内存根本不存在,但养成这个习惯后,以后处理百万级数据时,你会知道需要把数据放在文件或数据库中、在__getitem__里按需读取,而不是一股脑全加载到内存。

3. 从零实现一套可复现的训练流程

3.1 环境准备与依赖清单

动手之前,先确认环境。我建议使用 Python 3.9 或更高版本,PyTorch 2.0 及以上(2.x 的 API 和 1.x 基本兼容,但性能和新特性有优势)。另外需要 numpy、scikit-learn、matplotlib 这几个库。

安装命令很简单:

pip install torch scikit-learn matplotlib

如果你是 Windows 用户,想装 GPU 版 PyTorch,建议直接去 PyTorch 官网选对应的 CUDA 版本命令安装。CPU 版本跑这个演示也完全没问题。

环境就绪后,把随机种子固定下来,确保每次运行结果一致:

import random import numpy as np import torch def set_seed(seed=42): random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) torch.cuda.manual_seed_all(seed) set_seed(42)

3.2 自定义 Dataset 子类的完整写法

前面的通用子类可以用来包装任意特征矩阵和标签数组。加载 sklearn 内置数据并做划分和封装的完整代码如下:

from sklearn.datasets import load_diabetes from sklearn.model_selection import train_test_split from sklearn.preprocessing import StandardScaler from torch.utils.data import DataLoader # 1. 加载数据 diabetes = load_diabetes() X, y = diabetes.data, diabetes.target # 2. 划分训练集和测试集 X_train, X_test, y_train, y_test = train_test_split( X, y, test_size=0.2, random_state=42 ) # 3. 标准化 scaler = StandardScaler() X_train = scaler.fit_transform(X_train) X_test = scaler.transform(X_test) # 4. 封装成 Dataset train_dataset = TabularDataset(X_train, y_train) test_dataset = TabularDataset(X_test, y_test) # 5. 创建 DataLoader train_loader = DataLoader( train_dataset, batch_size=32, shuffle=True, num_workers=0 ) test_loader = DataLoader( test_dataset, batch_size=32, shuffle=False, num_workers=0 )

这里有一个计算可以顺手做一下:442 个样本,80% 是 353 个训练样本,20% 是 89 个测试样本。当 batch_size=32 时,353 个训练样本会被分成 11 个 batch,其中前 10 个 batch 各有 32 个样本,最后一个 batch 只有 33 个样本(353 - 10×32 = 33)。严格说是 11 个 batch 里最后一个 batch 是 33 个样本,因为 10×32=320,353-320=33。

测试集 89 个样本,batch_size=32 会产生 3 个 batch,最后一个是 25 个样本。最后一个 batch 不满,PyTorch 默认会把剩余样本凑成一个 batch,行为是自动的,不用手动处理。

3.3 模型搭建:从线性回归到两层 MLP

数据管道搭好后,接下来是模型。对于糖尿病数据集这类表格数据,最基础的方案是线性回归。如果用全连接网络来表达,就是一个不带激活函数的单层网络:输入 10 个特征,输出 1 个值。

但线性模型表达能力有限。现实中特征和目标之间往往存在非线性关系,简单堆一个线性层效果不会太好。我建议直接上一个小型多层感知机(MLP),两层隐藏层,每层 64 个神经元,激活函数用 ReLU:

import torch.nn as nn class DiabetesMLP(nn.Module): def __init__(self, in_features=10, hidden_dim=64): super().__init__() self.net = nn.Sequential( nn.Linear(in_features, hidden_dim), nn.ReLU(), nn.Linear(hidden_dim, hidden_dim), nn.ReLU(), nn.Linear(hidden_dim, 1) ) def forward(self, x): return self.net(x)

为什么是两层隐藏层?对于 442 个样本、10 个特征的数据,单层隐藏层已经能拟合不错的非线性关系,再加一层可以进一步提升表达能力,但要注意过拟合风险。如果你发现训练集 R² 远高于测试集 R²,就该考虑减少层数、加大正则化或提前停止训练了。这个数据集上用两层 64 维的隐藏层,配合适当的训练轮数,效果通常比较理想。

3.4 训练循环与评估细节

训练循环看似简单,但里面有不少容易被忽略的细节。完整代码如下:

import torch.optim as optim import torch.nn.functional as F from sklearn.metrics import r2_score, mean_squared_error device = torch.device("cuda" if torch.cuda.is_available() else "cpu") model = DiabetesMLP().to(device) criterion = nn.MSELoss() optimizer = optim.Adam(model.parameters(), lr=1e-3) n_epochs = 300 train_losses = [] test_r2s = [] for epoch in range(n_epochs): model.train() total_loss = 0.0 n_samples = 0 for X_batch, y_batch in train_loader: X_batch, y_batch = X_batch.to(device), y_batch.to(device) optimizer.zero_grad() pred = model(X_batch) loss = criterion(pred, y_batch) loss.backward() optimizer.step() total_loss += loss.item() * X_batch.size(0) n_samples += X_batch.size(0) avg_loss = total_loss / n_samples train_losses.append(avg_loss) # 每个 epoch 结束后在测试集上评估 model.eval() preds = [] y_true = [] with torch.no_grad(): for X_batch, y_batch in test_loader: X_batch, y_batch = X_batch.to(device), y_batch.to(device) pred_batch = model(X_batch) preds.append(pred_batch.cpu()) y_true.append(y_batch.cpu()) preds = torch.cat(preds, dim=0).numpy() y_true = torch.cat(y_true, dim=0).numpy() r2 = r2_score(y_true, preds) test_r2s.append(r2) if (epoch + 1) % 50 == 0: print(f"Epoch {epoch+1:03d} | Train Loss: {avg_loss:.4f} | Test R2: {r2:.4f}")

训练细节里,我特别想强调两点:

第一,model.train()和model.eval()要成对出现。虽然这个 MLP 里没有 dropout 和 batch normalization,但在复杂模型里,漏掉.eval()会直接导致评估结果异常。养成习惯后,你以后跑任何 PyTorch 模型都不会漏。

第二,评估时要包torch.no_grad()。这个上下文管理器告诉 PyTorch 不需要计算梯度图,可以大幅减少内存占用并加快推理速度。反向传播只在训练阶段需要,评估阶段关心的是前向结果。

第三,loss 累加时用loss.item() * X_batch.size(0),再除以总样本数,得到的是整个 epoch 的平均 loss,而不是每个 batch loss 的简单平均。因为最后一个 batch 样本数可能不满 32,如果直接对所有 batch loss 取平均,小 batch 的 loss 会被不公平地放大。

4. 训练效果评估与调优:光跑通不够,得会看指标

4.1 回归任务的评估指标:MSE、RMSE、R²

糖尿病患者数据集是回归任务,分类准确率在这里没有意义。常用的指标有三个:MSE、RMSE、R²。它们的侧重点各不相同。

MSE(均方误差):所有样本预测值与真实值之差的平方的平均数。它直接对应训练 loss,但对异常值比较敏感——一个离群样本的平方误差会占据很大权重。

RMSE(均方根误差):MSE 的算术平方根。它的优势是量纲和原始目标值一致。糖尿病数据集的目标值范围大致在 25 到 346 之间,如果 RMSE 在 50 左右,说明平均预测偏差约为 50 个单位,业务含义非常直观。

R²(决定系数):最常用的回归评估指标,表示模型解释了目标值多大比例的方差。R² = 1 是完美预测,R² = 0 意味着模型效果等同于直接预测平均值,R² 为负时说明模型比瞎猜还差。

一个模型在糖尿病数据集上的典型表现是:R² 大约在 0.45 到 0.55 之间,RMSE 大约在 50 到 60 之间。这不是模型不行,而是这份数据本身的特征信息有限,仅靠 10 个临床指标很难精确预测病情进展的数值。看到这个结果不用失望,R² 超过 0.5 在这个数据集上已经算不错的成绩了。

4.2 学习率、batch size、epoch 的参数组合实验

调参最好一次只改一个变量,否则出了问题你不知道是哪个参数导致的。我整理了一张参数组合参考表,都是在这个数据集上实测过的结果:

配置学习率batch sizeepoch测试 R² 表现
A1e-232300约 0.40~0.45,训练震荡较大
B1e-332300约 0.47~0.53,稳定收敛
C1e-432300约 0.45~0.50,收敛偏慢
D1e-364300约 0.45~0.50,略低于 batch 32
E1e-332600约 0.48~0.54,与 B 接近

学习率 1e-2 在这个模型上明显偏高,训练曲线像过山车;1e-3 表现稳定,是稳妥的起点;1e-4 偏保守,收敛太慢,300 轮可能不够。batch size 从 32 调到 64 效果略降,因为参数更新次数减半了。epoch 从 300 增加到 600,R² 提升有限,说明模型已经接近该结构的天花板。

我的建议是:先用学习率 1e-3、batch size 32 跑 300 轮作为基线,再根据曲线决定下一步动作。如果训练 loss 还在持续下降,就加大 epoch;如果已经平稳,就考虑调整模型结构而非继续训练。

4.3 过拟合的识别与缓解

小数据集上训练 MLP,过拟合几乎是必然要面对的问题。识别过拟合的方法很简单:训练集的 R² 远高于测试集。比如训练集 R² 达到 0.7,测试集只有 0.4,这中间 0.3 的差距就是过拟合的代价。

缓解手段从易到难排序:

  • 提前停止(early stopping):训练时监测测试集指标,一旦连续多个 epoch 不再提升就停止训练。
  • 减小模型容量:把隐藏层维度从 64 降到 32,或者去掉一层隐藏层。
  • 加权重衰减(weight decay):在优化器中设置weight_decay=1e-4,相当于对权重做 L2 正则化,限制模型复杂度。
  • 加 Dropout:在隐藏层后插入 Dropout 层,强迫模型不依赖少数特征。

对于这个数据集,我实测最有效的是提前停止和权重衰减。Dropout 在小数据集上容易让模型欠拟合,需要配更高的 epoch 数才有效果。不管用哪种方法,核心思路都是:在不过度牺牲训练集表现的前提下,尽量缩小训练集和测试集的差距。

5. 常见问题与排查技巧实录

5.1 Windows 下 num_workers 的经典报错

在 Windows 上使用num_workers>0且代码没有放在主模块保护内,运行时会报RuntimeError: An attempt has been called in the main process...,这是新手遇到频率极高的问题。解决办法是在脚本入口处加上:

if __name__ == "__main__": main()

把训练逻辑放进main()函数里。这是因为 Windows 下多进程加载数据需要重新导入主模块,没有if __name__ == "__main__"保护会递归创建进程导致崩溃。如果只是跑实验脚本,直接把num_workers设为 0 更省心。

5.2 pin_memory 的收益与使用条件

pin_memory=True是 DataLoader 的又一个参数。它告诉 PyTorch 将加载的数据张量放入锁页内存,当数据要转移到 GPU 时,可以显著加快传输速度。但有两个使用条件:第一,你的数据最终要拷贝到 GPU;第二,使用tensor.to(device, non_blocking=True)才能真正享受加速。如果全程 CPU 训练,这个参数没有任何收益,甚至可能带来轻微性能下降。

5.3 训练不收敛的排查清单

训练 loss 不下降或直接变成 NaN,这类问题我给自己总结了一张排查清单:

  1. 数据预处理检查:特征是否标准化?目标值是否异常大?如果目标值动辄上千,MSE loss 会非常大,可能淹没梯度信号。
  2. 学习率检查:学习率太高会导致 loss 震荡甚至爆炸,太低会导致几乎不更新。出现 NaN 时优先调低学习率。
  3. 标签形状检查:标签形状是(batch_size,)还是(batch_size, 1)?两者混用可能导致广播维度错乱,尤其在计算 MSE 时,PyTorch 有时不会报错,但结果完全不对。
  4. 模型输出检查:模型最后输出维度是否为 1?如果是分类任务的输出头忘了改,维度对不上直接报错。

5.4 一个隐秘的坑:从 numpy 到 tensor 的 dtype 问题

sklearn 默认返回 float64 类型的数组,而 PyTorch 默认权重是 float32。直接对 float64 的 numpy 数组执行torch.tensor(arr),默认会保留 float64。模型权重是 float32,输入是 float64,虽然 PyTorch 有时会自动转换,但某些算子会报 dtype mismatch 错误,或者悄无声息地多一次内存拷贝。

稳妥的做法是在构造 Dataset 时显式指定 dtype:

self.X = torch.tensor(X, dtype=torch.float32)

这行代码看似不起眼,但它能避免一大类隐蔽的运行时错误。我自己就遇到过:代码在 CPU 上跑得好好的,一到 GPU 就报 dtype 不匹配,查了半天才发现是这里的问题。

5.5 复现实验的最后一公里:保存与加载

实验跑通后,别忘了把模型参数保存下来。PyTorch 的推荐做法是只保存状态字典,不保存整个模型对象:

torch.save(model.state_dict(), "diabetes_mlp.ckpt") # 加载方式 model = DiabetesMLP() model.load_state_dict(torch.load("diabetes_mlp.ckpt", weights_only=True)) model.eval()

weights_only=True是 PyTorch 2.x 增加的安全选项,防止加载恶意构造的 pickle 文件。养成这个习惯,以后在正式项目和团队协作中会少很多麻烦。

最后分享几个提升效率的小习惯

前面都是操作层面的事,最后聊点方法论。我在实践里有一个习惯:每个小项目跑通后,把最终的模型配置和关键结果记成一个文本文件,和训练脚本放在同一目录。换环境、换设备、隔几个月再回来看项目时,这份记录比任何注释都管用。另一个习惯是调试数据管道时,先取一个 batch 的数据打印出来看形状、dtype 和数值范围,确认没问题再正式训练。这一步一分钟都不到,但能省下后面排查烂摊子的几个小时。

这一套“DataLoader 加载糖尿病数据集训练”的流程走下来,你实际上已经掌握了一个可以迁移到大部分表格数据任务的完整模板。换数据、换模型、换指标,但数据加载和训练循环的骨架基本不变。后面再遇到 Kaggle 的房价预测、股票走势拟合、工业参数回归这类问题,都可以用同一套思路快速上手。至于更高级的特性,比如自定义采样器、多进程流式加载、超大规模数据集的分布式分片,都是在这条主干上长出来的枝叶,基础打牢了,那些自然不怵。

需要专业的网站建设服务?

联系我们获取免费的网站建设咨询和方案报价,让我们帮助您实现业务目标

立即咨询