1. 从一次数据加载卡顿说起：为什么需要搞懂这三个类
刚接触深度学习那会儿，我写过一个图像分类的小项目。数据集不大，几千张图，我图省事，把所有图片一次性读进内存，转成张量，然后手动写了个循环，每个batch切一刀喂给模型。训练倒是能跑，但每次改个batch size或者想加个打乱顺序的功能，就得回去改循环里的逻辑，改来改去代码越来越乱。更难受的是，后来换了个稍大一点的数据集，内存直接爆了，程序还没开始训练就崩了。
那次之后我才认真去看了PyTorch的数据加载体系，发现TensorDataset、Dataset、DataLoader这三个东西其实是一套组合拳，各司其职。Dataset是抽象基类，定义了"数据从哪来、怎么取一条"的规范；TensorDataset是官方给的一个现成实现，专门用来包装已经转成张量的数据；DataLoader则负责调度，把Dataset里的数据按batch取出来、打乱顺序、并行加载。三者配合起来，代码干净，内存可控，扩展也方便。
这篇内容就是把我踩过的坑和后来总结的用法梳理一遍。不管你是刚学PyTorch的新手，还是写过几个项目但一直用土办法加载数据的人，看完应该能把这套体系用顺。我会从最基础的概念讲起，然后给出可直接复现的代码，再聊几个实际项目中容易忽略的细节，最后说说什么时候该自己写Dataset、什么时候用TensorDataset就够了。
2. 三个角色的分工：Dataset是规范，TensorDataset是快捷实现，DataLoader是调度器
2.1 Dataset到底定义了什么
PyTorch里的Dataset是一个抽象类，它本身不能直接实例化使用，作用是给所有数据集定义一个统一的接口。你继承它之后，必须实现两个方法：__len__和__getitem__。
__len__返回数据集里有多少条样本，DataLoader靠它知道一共要取多少数据、一个epoch有多少个batch。__getitem__接收一个索引，返回一条样本，这条样本可以是张量、可以是元组、也可以是字典，具体格式由你自己定。
为什么要有这么一层抽象？因为实际项目里数据来源五花八门：有的是硬盘上的图片文件夹，有的是CSV文件，有的是数据库查询结果，还有的是实时生成的合成数据。如果每个地方都自己写加载逻辑，代码风格会非常混乱。有了Dataset这个统一接口，上层的数据调度逻辑就不用关心底层数据长什么样，只管按索引取就行。
我个人的理解是，Dataset解决的是"数据怎么读"的问题，它把读取逻辑封装起来，让训练循环变得干净。你想想，如果训练代码里混着文件路径拼接、图像解码、归一化这些操作，那代码根本没法维护。
2.2 TensorDataset帮你省掉重复劳动
TensorDataset是Dataset的一个子类，官方已经帮你实现好了__len__和__getitem__。它的用法很简单：把几个张量传进去，它会按第一个维度对齐，每次取索引时返回这些张量在对应位置的切片。
举个例子，你有特征张量X，形状是(1000, 20)，标签张量y，形状是(1000,)。用TensorDataset(X, y)包一下，取第i条时就会返回(X[i], y[i])。就这么直接。
它适合什么场景？数据已经全部在内存里、已经转成张量、不需要复杂的预处理。比如你从numpy数组转过来的数据、从预训练特征提取器输出的特征、或者做小规模实验时手动构造的数据。这种情况下用TensorDataset最省事，不用自己写类。
但它也有明显的局限。第一，它要求所有数据都能放进内存，数据量大了就不行。第二，它不支持自定义的预处理逻辑，比如你想在取数据时做随机裁剪、加噪声，TensorDataset做不到。第三，它返回的是固定顺序的元组，如果你想要字典格式的输出，得自己再包一层。
2.3 DataLoader才是真正干活的那个
DataLoader接收一个Dataset对象，然后负责：按batch size把数据分组、决定是否打乱顺序、决定用几个进程加载、决定怎么把多条样本拼成一个batch。
它内部的工作流程大致是这样的：先根据__len__知道总共有多少条，然后按batch_size算出有多少个batch。如果shuffle=True，就生成一个随机排列的索引列表。然后按索引去Dataset里取数据，取够一个batch就通过collate_fn拼起来返回。
collate_fn是个容易被忽略但很重要的参数。默认情况下，它会把一个batch里的样本按位置堆叠成张量。比如每条样本是(特征, 标签)，一个batch有32条，默认就会返回两个张量，形状分别是(32, 特征维度)和(32,)。但如果你的样本长度不一致，比如文本序列，默认的collate就会报错，这时候就需要自定义collate_fn来做padding。
还有一个参数num_workers，控制用几个子进程加载数据。设成0表示在主进程里加载，设成大于0会开子进程。这个参数对训练速度影响很大，但也不是越大越好，后面会细说。
3. 从零写一个可运行的例子：TensorDataset配DataLoader的完整流程
3.1 构造数据与包装
先来看一段可以直接跑的代码。假设我们做一个回归任务，输入是20维特征，输出是一个连续值。
import torch from torch.utils.data import TensorDataset, DataLoader # 构造模拟数据 torch.manual_seed(42) N = 1000 X = torch.randn(N, 20) true_w = torch.randn(20, 1) y = X @ true_w + 0.1 * torch.randn(N, 1) # 包装成TensorDataset dataset = TensorDataset(X, y) print(len(dataset)) # 1000 print(dataset[0]) # (X[0], y[0]) 两个张量这里dataset[0]返回的是一个元组，第一个元素是20维特征，第二个元素是1维标签。注意标签的形状是(1,)而不是标量，因为原始y是(N, 1)，切片后保持了这个维度。
3.2 DataLoader的配置与迭代
dataloader = DataLoader( dataset, batch_size=32, shuffle=True, num_workers=0, drop_last=False ) for batch_X, batch_y in dataloader: print(batch_X.shape, batch_y.shape) break # 输出: torch.Size([32, 20]) torch.Size([32, 1])这段代码里几个参数值得说一下。batch_size=32表示每个batch取32条。shuffle=True表示每个epoch开始前重新打乱索引顺序，这对训练很重要，能防止模型学到样本顺序带来的虚假规律。drop_last=False表示如果最后一批不够32条也保留，如果设成True就会丢掉最后那批不足的。
num_workers=0表示不用子进程，数据在主进程里加载。对于这种全在内存里的张量数据，设成0反而更快，因为开子进程有额外的进程创建和数据传输开销。
3.3 训练循环里怎么用
import torch.nn as nn model = nn.Sequential( nn.Linear(20, 64), nn.ReLU(), nn.Linear(64, 1) ) optimizer = torch.optim.Adam(model.parameters(), lr=1e-3) loss_fn = nn.MSELoss() for epoch in range(5): model.train() total_loss = 0.0 for batch_X, batch_y in dataloader: optimizer.zero_grad() pred = model(batch_X) loss = loss_fn(pred, batch_y) loss.backward() optimizer.step() total_loss += loss.item() * batch_X.size(0) print(f"epoch {epoch}, avg loss {total_loss / len(dataset):.4f}")这个训练循环就是标准写法。每次从dataloader里取出一个batch，前向、算损失、反向、更新。注意loss.item() * batch_X.size(0)是为了按样本数加权平均，因为最后一个batch可能不满。
跑下来你会发现，用DataLoader之后，训练代码里完全没有数据加载的细节，全是模型相关的逻辑。这就是这套体系的价值：关注点分离。
4. 当TensorDataset不够用：自定义Dataset的典型场景与写法
4.1 什么时候必须自己写Dataset
TensorDataset虽然方便，但下面这些情况它搞不定：
- 数据在硬盘上，不能一次性读进内存，比如几万张高分辨率图片。
- 每条样本需要不同的预处理，比如随机裁剪、随机翻转、加噪声。
- 数据来源是数据库或网络接口，需要按需查询。
- 样本长度不一致，需要动态padding。
- 需要返回字典格式而不是元组，方便后续处理。
这些场景下就得继承Dataset自己实现。核心还是那两个方法，但__getitem__里可以写任意逻辑。
4.2 一个图像数据集的完整实现
假设我们有一个图片文件夹，每张图对应一个类别标签，标签存在一个CSV文件里。下面是一个简化但可运行的实现：
import os import csv import torch from torch.utils.data import Dataset from PIL import Image import torchvision.transforms as T class ImageFolderDataset(Dataset): def __init__(self, root_dir, label_file, transform=None): self.root_dir = root_dir self.transform = transform self.samples = [] with open(label_file, 'r') as f: reader = csv.reader(f) next(reader) # 跳过表头 for row in reader: img_name, label = row[0], int(row[1]) self.samples.append((img_name, label)) def __len__(self): return len(self.samples) def __getitem__(self, idx): img_name, label = self.samples[idx] img_path = os.path.join(self.root_dir, img_name) image = Image.open(img_path).convert('RGB') if self.transform: image = self.transform(image) return image, label这个类里，__init__只做轻量的事：读标签文件、存路径列表。真正的图片读取和解码放在__getitem__里，这样只有被取到的样本才会加载，内存占用可控。
transform参数用来传入预处理管道，比如：
train_transform = T.Compose([ T.RandomResizedCrop(224), T.RandomHorizontalFlip(), T.ToTensor(), T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) ])注意ToTensor()必须放在归一化之前，因为Normalize只接受张量。这个顺序搞反了会报错，我见过不止一个新手在这里卡住。
4.3 自定义Dataset配合DataLoader的注意事项
自定义Dataset和DataLoader配合时，有几个点容易出问题。
第一，__getitem__里不要修改类的成员变量。因为如果num_workers > 0，每个子进程会复制一份Dataset对象，你在子进程里改的东西主进程看不到，而且多个worker之间也不共享。所有需要随机的操作，比如随机裁剪，应该用transform里的随机函数，它们内部有自己的随机状态。
第二，如果__getitem__里有可能抛异常，比如图片损坏、文件缺失，最好在内部捕获并返回一个替代样本，否则整个训练会崩。我一般会加个try-except，遇到坏样本就返回一个全零张量和忽略标签，或者直接随机换一个索引重取。
第三，__len__返回的长度必须和实际可取的索引范围一致。如果返回的长度比实际样本多，DataLoader取到后面会报IndexError。
5. DataLoader那些影响训练效率和结果的参数
5.1 num_workers设多少合适
num_workers是最常被问到的参数。它的作用是开多少个子进程来并行加载数据。设成0表示不用子进程，在主进程里串行加载。
什么时候该设大于0？当数据加载成为瓶颈的时候。比如你从硬盘读图片、做复杂的图像增强，这些操作耗时较长，如果串行做，GPU就会闲着等数据。这时候开几个worker并行加载，能让GPU一直有活干。
但也不是越大越好。每个worker都是一个独立的进程，会占用内存和CPU。worker太多，进程间切换的开销反而拖慢速度。经验值是设成CPU核心数的一半左右，比如8核机器设4。但最终还是要实测，用time模块测一下一个epoch的耗时，对比不同设置。
还有一个坑：在Windows上，num_workers > 0时，DataLoader的构造必须放在if __name__ == '__main__':保护块里，否则会无限递归创建进程。这个在Linux上不是问题，但跨平台代码要注意。
5.2 shuffle、drop_last与sampler的关系
shuffle=True本质上是在每个epoch开始时，用RandomSampler替换默认的SequentialSampler。如果你传了自定义的sampler参数，那shuffle就必须是False，否则会报冲突。
drop_last决定是否丢掉最后一个不满的batch。什么时候该设True？当你的模型里有BatchNorm层，而最后一个batch只有1条样本时，BatchNorm会报错，因为它在训练时需要至少2条样本来计算方差。这种情况下设drop_last=True能避免报错。
但丢掉最后一批也意味着损失了一部分数据。如果数据量本来就少，丢掉可能影响训练效果。折中办法是把batch size调小一点，让最后一批不至于只有1条。
5.3 collate_fn的默认行为与自定义时机
默认的collate_fn会把一个batch里的样本按结构堆叠。如果每条样本是(张量A, 张量B)，它会返回(堆叠后的A, 堆叠后的B)。堆叠要求所有样本在对应位置上的形状一致。
什么时候需要自定义？最常见的是变长序列。比如每条样本是一个不同长度的文本序列，默认collate会因为形状不一致而报错。这时候需要写一个函数，把序列padding到batch内最大长度。
def pad_collate(batch): sequences, labels = zip(*batch) max_len = max(seq.size(0) for seq in sequences) padded = torch.zeros(len(sequences), max_len, dtype=torch.long) for i, seq in enumerate(sequences): padded[i, :seq.size(0)] = seq return padded, torch.tensor(labels)这个函数接收一个列表，每个元素是__getitem__返回的样本。先解包，算出最大长度，然后创建全零张量，把每条序列填进去。最后返回padding后的张量和标签张量。
自定义collate_fn时要注意，它是在主进程里执行的，不在worker里。所以如果padding逻辑很重，会成为瓶颈。一般padding操作本身很快，问题不大。
6. 几个我实际踩过的坑和对应的排查思路
6.1 内存暴涨：worker进程的数据复制
有一次我做一个较大的数据集，num_workers设了8，结果内存直接爆了。排查后发现，每个worker进程都会复制一份Dataset对象。如果Dataset的__init__里加载了大量数据到内存，那8个worker就是8份，内存直接乘以8。
解决办法是把重数据的加载推迟到__getitem__里，__init__只存路径和元信息。这样每个worker只在需要时才读具体样本，不会一开始就复制全部数据。
另一个相关的问题是，如果__getitem__里缓存了数据，比如用一个字典存已读过的图片，那每个worker都有自己的缓存，主进程看不到，也起不到共享作用。这种缓存只在单个worker内部有效，跨worker不共享。
6.2 训练结果不可复现：随机种子的设置
用DataLoader时，如果shuffle=True且num_workers > 0，即使你设了torch.manual_seed，每次运行的结果也可能不一样。原因是每个worker有自己的随机状态，主进程的种子不会自动传给它们。
要完全复现，需要给worker设置种子。PyTorch提供了一个worker_init_fn参数：
def worker_init_fn(worker_id): np.random.seed(42 + worker_id) torch.manual_seed(42 + worker_id) dataloader = DataLoader(dataset, batch_size=32, shuffle=True, num_workers=4, worker_init_fn=worker_init_fn)这样每个worker的随机种子是确定的，配合主进程的种子，就能复现。不过要注意，即使这样，如果用了CUDA的某些非确定性操作，还是可能不完全一致。完全复现是个系统工程，这里只是解决数据加载部分的随机性。
6.3 最后一个batch引发的形状错误
前面提过BatchNorm在单样本batch时会报错。还有一种情况是模型里用了view或reshape，而最后一个batch的样本数不是预期值，导致形状对不上。
比如模型里写了x.view(batch_size, -1)，但最后一个batch实际只有16条而batch_size设的是32，就会报错。正确写法是用x.view(x.size(0), -1)，让形状根据实际数据来。
这个坑的本质是：不要假设每个batch的大小都等于设定的batch_size。要么用drop_last=True保证，要么在代码里用动态形状。
6.4 数据加载成为瓶颈的识别方法
怎么判断数据加载是不是瓶颈？一个简单办法是在训练循环里计时。如果GPU利用率低，而CPU占用高，大概率是数据加载拖后腿。
更精确的做法是分别测数据加载时间和模型前向反向时间：
import time # 测数据加载 t0 = time.time() for _ in dataloader: pass t_data = time.time() - t0 # 测模型计算 batch = next(iter(dataloader)) t0 = time.time() for _ in range(100): pred = model(batch[0]) loss = loss_fn(pred, batch[1]) loss.backward() t_model = time.time() - t0如果t_data和t_model接近甚至更大，说明数据加载是瓶颈，可以考虑增加num_workers、简化预处理、或者把数据预处理提前做好存成缓存文件。
7. 从TensorDataset到自定义Dataset的平滑过渡建议
如果你现在用的是TensorDataset，但发现数据量变大或者需要加预处理，想迁移到自定义Dataset，我建议按这个顺序来。
第一步，先把TensorDataset替换成一个最简单的自定义Dataset，__getitem__里直接返回原来TensorDataset返回的东西。跑一遍确认训练结果和之前一致，说明接口对接没问题。
第二步，把数据加载逻辑从__init__里挪到__getitem__里。比如原来是__init__里读所有图片，现在改成__init__只存路径，__getitem__里读单张。这一步做完，内存占用应该会明显下降。
第三步，加入预处理逻辑。把transform管道加到__getitem__里，注意训练和验证要用不同的transform。验证集一般不做随机增强，只做resize和归一化。
第四步，调整DataLoader参数。数据从内存变成硬盘读取后，加载耗时增加，这时候把num_workers调大，观察训练速度变化。
这个渐进式迁移的好处是每一步都可验证，出了问题容易定位是哪一步引入的。直接重写整个数据管道然后调试，往往会被各种小问题缠住。
另外说一个我个人的习惯：不管用TensorDataset还是自定义Dataset，我都会先写一个小的测试脚本，单独把Dataset和DataLoader拿出来跑一遍，打印几个batch的形状和内容，确认没问题了再接到训练循环里。这样能把数据问题和模型问题分开排查，省很多时间。
最后提一下，PyTorch还有IterableDataset这个类，适合流式数据或者数据量大到无法用索引访问的场景。但那是另一个话题了，对于大多数项目，Dataset加DataLoader这套组合已经够用。把这两个用熟，再根据实际需求决定要不要上更复杂的方案。
