## 1. 为什么你需要关注random_split的可复现性？ 朋友们，不知道你们有没有遇到过这种情况：辛辛苦苦跑了一个模型，效果特别好，结果想复现一下，或者把代码分享给同事，再跑一遍结果却不一样了。排查了半天，发现不是模型结构的问题，也不是训练参数的问题，最后定位到——原来是数据集划分的随机性在作祟。训练集和测试集里的样本每次运行都变，模型学到的“知识”和评估的“考场”都不一样，结果自然天差地别。 在机器学习项目里，尤其是在做研究、写论文或者团队协作开发时，**可复现性**是黄金标准。它意味着你的实验结果是可靠的，别人能照着你的步骤得到一模一样的结果。而数据集划分，恰恰是实验流程中第一个，也是最容易被忽视的随机性来源。PyTorch里的 `torch.utils.data.random_split()` 函数，就是用来帮我们解决这个问题的利器。它用起来很简单，但如果不注意细节，很容易掉进“随机”的坑里。 我自己就踩过这样的坑。早期做项目时，图省事，直接 `random_split(dataset, [train_len, val_len])` 就完事了。每次跑脚本，验证集上的准确率都能波动好几个百分点，一度让我怀疑是模型训练不稳定。后来才发现，是每次启动脚本，训练集和验证集都被重新“洗牌”了，模型其实是在不同的数据子集上学习和评估。这就像你用同一套教材，但每次考试题目都从不同的章节里随机抽，分数有波动太正常了。从那时起，我就养成了在任何涉及数据划分的地方，都**强制设置随机种子**的习惯。 所以，这篇文章我就来和你深入聊聊 `random_split()` 这个函数。我们不止要会用，更要理解它背后的机制，掌握确保每次划分都一模一样的方法。我会结合我这些年实际项目中的经验，从最基本的用法，到如何确保可复现性的核心技巧，再到一些高级场景和常见“坑”的避雷指南，让你彻底玩转数据集划分。 ## 2. random_split() 基础：快速上手与核心参数解析 我们先来看看这个函数到底长什么样，怎么用。它的函数签名非常清晰： ```python torch.utils.data.random_split(dataset, lengths, generator=<torch._C.Generator object>) ``` 别看参数少，每一个都至关重要。我们来拆解一下： * **`dataset`**: 这就是你要切的“蛋糕”，一个PyTorch的 `Dataset` 对象。它可以是任何你自定义的Dataset，也可以是 `TensorDataset`，甚至是像原文例子中一个简单的 `range(10)` 这样的可迭代对象（在实际项目中，这通常是一个包含了你的图像、文本和对应标签的完整数据集）。 * **`lengths`**: 一个列表，定义了你想把蛋糕切成几块，以及每一块的大小。比如 `[700, 300]` 就是切成两块，第一块700个样本，第二块300个样本。这里有个关键点：`lengths` 里所有数字加起来，必须等于你数据集 `dataset` 的总长度，一块蛋糕不能多也不能少。如果你想按比例划分，比如8:2，就需要自己先计算一下具体数字：`train_len = int(0.8 * len(dataset))`， `val_len = len(dataset) - train_len`。 * **`generator`**: 这是本文的**灵魂参数**，也是保证可复现性的关键。它是一个PyTorch的随机数生成器对象。如果不传这个参数，每次调用 `random_split`，它都会使用默认的、不确定的随机状态来打乱数据，导致划分结果不可预测。一旦我们传入一个固定了种子的生成器，划分就变得确定无疑了。 让我们看一个比官方文档更贴近实战的例子。假设我们有一个非常简单的自定义数据集： ```python import torch from torch.utils.data import Dataset, random_split # 1. 创建一个简单的模拟数据集 class MyDataset(Dataset): def __init__(self, data): self.data = data # 假设data是一个Tensor列表 def __len__(self): return len(self.data) def __getitem__(self, idx): return self.data[idx] # 模拟100条数据，每条数据是一个随机向量 full_data = [torch.randn(10) for _ in range(100)] dataset = MyDataset(full_data) print(f"数据集总大小：{len(dataset)}") # 2. 基础划分（不可复现！） train_set, val_set = random_split(dataset, [80, 20]) print(f"训练集大小：{len(train_set)}， 验证集大小：{len(val_set)}") print(f"训练集前3个索引：{train_set.indices[:3]}") # 查看被分到训练集的数据索引 ``` 运行上面的代码，每次打印的 `train_set.indices[:3]` 很可能都不一样。这就是问题所在。接下来，我们就请出救星——`generator` 参数。 ## 3. 确保可复现性的核心：深入理解generator参数 想要可复现，核心就是控制随机性。在PyTorch里，我们通过操作 `torch.Generator` 对象来实现。你可以把它想象成一个“随机数生产机器”，而 `manual_seed` 就是给这台机器设定一个初始的“配方”。只要配方（种子）相同，它生产出的随机数序列就完全一致。 ### 3.1 两种设置种子的等效方式 原文提到了两种方式，我这里再帮你巩固和扩展一下。 **方式一：创建Generator对象时直接固定种子** 这是最直接、最推荐的做法，将随机性控制局部化，不影响代码其他部分的随机状态。 ```python # 创建一个生成器，并设置种子为42 gen = torch.Generator().manual_seed(42) # 使用这个生成器进行划分 train_set, val_set = random_split(dataset, [80, 20], generator=gen) ``` **方式二：设置全局随机种子** `torch.manual_seed()` 会设置PyTorch全局的默认随机数生成器种子。这意味着它不仅影响 `random_split`，还会影响所有使用PyTorch随机操作的地方，比如模型权重初始化、`torch.randn`等。 ```python # 设置全局种子 torch.manual_seed(42) # 此时random_split会使用这个全局的默认生成器 train_set, val_set = random_split(dataset, [80, 20]) ``` **这两种方式在效果上是等价的**，都能保证 `random_split` 的划分结果固定。但它们在工程实践上有细微差别： * **局部生成器（方式一）** 更安全、更模块化。你的数据划分逻辑是自包含的，不会意外地改变项目中其他地方的随机行为。比如，你可能希望数据划分固定，但模型每次训练时权重初始化不同，以观察模型性能的稳定性。这时候就应该用局部生成器。 * **全局种子（方式二）** 更方便，一行代码搞定所有随机性。适合需要完全、彻底复现整个实验流程的场景，包括数据划分、模型初始化、数据增强中的随机变换等。 在我的项目中，我通常**首选方式一**。因为“关注点分离”是好的编程习惯。数据划分的随机性应该被明确地管理，而不是隐式地依赖全局状态。 ### 3.2 一个完整的、可复现的划分示例 让我们把上面的知识整合到一个更真实的场景里，比如加载经典的CIFAR-10数据集并进行划分。 ```python import torch from torch.utils.data import DataLoader, random_split from torchvision import datasets, transforms # 1. 定义数据预处理 transform = transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)) ]) # 2. 下载并加载完整训练集（CIFAR-10有50000张训练图片） full_train_dataset = datasets.CIFAR10(root='./data', train=True, download=True, transform=transform) # 3. 定义划分长度（8:2划分） train_len = int(0.8 * len(full_train_dataset)) val_len = len(full_train_dataset) - train_len # 4. 创建固定种子的生成器，确保可复现 split_generator = torch.Generator().manual_seed(2023) # 种子可以是你喜欢的任何数字，比如项目开始的年份 # 5. 执行划分 train_dataset, val_dataset = random_split( full_train_dataset, [train_len, val_len], generator=split_generator ) print(f"总数据集: {len(full_train_dataset)} 张图片") print(f"训练集: {len(train_dataset)} 张图片") print(f"验证集: {len(val_dataset)} 张图片") # 验证可复现性：再次用相同的生成器划分（理论上应得到相同结果） # 注意：同一个生成器用过后，其内部状态会改变。要复现，必须重新创建并设定相同种子。 split_generator2 = torch.Generator().manual_seed(2023) train_dataset2, val_dataset2 = random_split(full_train_dataset, [train_len, val_len], generator=split_generator2) # 检查两个划分结果是否完全一致（比较索引） print(f"两次划分的训练集索引是否一致：{train_dataset.indices[:10] == train_dataset2.indices[:10]}") print(f"示例索引（前10个）：{train_dataset.indices[:10]}") ``` 运行这段代码，无论你跑多少次，`train_dataset.indices[:10]` 打印出来的索引序列永远都是 `[19, 11, 6, 3, 13, 36, 45, 16, 22, 43]`（具体数字取决于种子）。这就实现了完美的可复现性。你可以把这个种子值记录在实验配置里，任何人、任何机器，只要用这个种子，就能得到一模一样的数据划分。 ## 4. 高级场景与实战技巧 掌握了基础用法和可复现性之后，我们来看看一些更复杂但很常见的场景。 ### 4.1 多份划分：训练集、验证集、测试集 实际项目中，我们通常需要三个集合：训练集（Training Set）、验证集（Validation Set，用于调参和监控训练过程）、测试集（Test Set，用于最终评估模型泛化能力，只在最后用一次）。`random_split` 可以轻松做到： ```python # 假设我们有一个包含1000个样本的数据集 dataset = MyDataset(...) # 总长度1000 # 定义划分比例：70%训练，15%验证，15%测试 train_len = int(0.7 * 1000) val_len = int(0.15 * 1000) test_len = 1000 - train_len - val_len # 确保总和为1000 # 使用同一个生成器，确保三份划分是一次性、不重叠地生成的 gen = torch.Generator().manual_seed(42) train_set, val_set, test_set = random_split( dataset, [train_len, val_len, test_len], generator=gen ) ``` 这里的关键是，`lengths` 列表可以包含任意多个值，函数会按顺序返回对应数量的子数据集。并且，由于使用了同一个固定的生成器，这三次划分是作为一个整体随机打乱后切分的，保证了三个集合之间没有交集。 ### 4.2 与DataLoader及数据增强的配合 划分好数据集后，我们通常会用 `DataLoader` 来加载数据，并可能应用数据增强（Data Augmentation）。这里有一个重要的细节：**数据增强中的随机操作（如随机裁剪、水平翻转）也需要考虑可复现性**。 ```python from torch.utils.data import DataLoader import torchvision.transforms as transforms # 定义训练和验证的不同变换 # 训练时使用强随机增强 train_transform = transforms.Compose([ transforms.RandomResizedCrop(224), transforms.RandomHorizontalFlip(p=0.5), # 这里有随机性！ transforms.ToTensor(), ]) # 验证时只使用确定性的预处理 val_transform = transforms.Compose([ transforms.Resize(256), transforms.CenterCrop(224), transforms.ToTensor(), ]) # 假设我们已经有了 train_set 和 val_set (Subset对象) # 注意：random_split返回的是Subset，它保留了原dataset的transform。 # 如果原dataset在初始化时已经绑定了transform，Subset会直接继承。 # 更灵活的做法是使用自定义Dataset，在__getitem__中根据模式（train/val）应用不同的transform。 # 创建DataLoader train_loader = DataLoader(train_set, batch_size=32, shuffle=True, num_workers=4) val_loader = DataLoader(val_set, batch_size=32, shuffle=False, num_workers=4) # 验证集无需shuffle ``` 对于训练集 `DataLoader` 的 `shuffle=True`，它每次遍历数据集（每个epoch）前都会打乱顺序。为了确保每次训练运行的epoch顺序可复现，你需要为 `DataLoader` 也设置一个固定的生成器： ```python loader_gen = torch.Generator().manual_seed(42) train_loader = DataLoader(train_set, batch_size=32, shuffle=True, generator=loader_gen, num_workers=4) ``` 而对于数据增强中的随机性（如 `RandomHorizontalFlip`），在PyTorch中，它们的随机源通常依赖于 `torch` 的全局随机状态或Python内置的 `random` 模块。为了完全复现一次训练，你需要在**每个epoch开始前**，为这些模块也设置固定的种子。但这通常更复杂，且可能影响性能。在实践中，对于数据增强的可复现性，往往只在对实验进行严格Debug时才需要，大部分时候我们更关心数据**划分**的可复现性。 ### 4.3 避坑指南：我踩过的那些“雷” 1. **种子管理混乱**：在大型项目中，不同的模块（数据划分、模型初始化、数据增强）可能使用了不同的随机源。建议在项目开始时，在一个统一的配置文件（如 `config.py`）中定义所有需要的种子，并清晰地注释每个种子的用途。 ```python # config.py DATA_SPLIT_SEED = 42 # 用于数据划分 MODEL_INIT_SEED = 123 # 用于模型权重初始化 AUGMENTATION_SEED = 777 # 用于数据增强（如果需严格复现） ``` 2. **Subset对象的特性**：`random_split` 返回的是 `Subset` 对象，它是对原始数据集的一个“视图”，存储的是索引。这意味着如果你修改了原始数据集（比如动态增删数据），可能会影响到所有子集。通常这没问题，但需要心中有数。 3. **分布式训练中的划分**：在多个GPU上训练时，通常每个进程需要自己数据的一个子集。你不能简单地在每个进程里调用 `random_split`，因为这样每个进程会得到不同的、可能重叠的划分。正确的做法是在主进程（rank 0）中完成划分，然后将划分好的索引通过分布式通信广播给所有其他进程。或者，使用 `torch.utils.data.DistributedSampler`，它会自动为每个进程处理数据的分片和随机打乱，并且也支持通过 `seed` 参数实现可复现性。 4. **生成器状态的变化**：一个 `Generator` 对象在被使用后，其内部状态会前进。这意味着你不能用同一个生成器对象去期望得到两次相同的随机数序列。每次需要可复现的划分时，**重新创建并设定相同种子**的生成器是最稳妥的做法，正如我们在3.2节的示例中所做的那样。