如何使用 Python 将数据集拆分为批次

在本指南中，您将学习：

• 什么是 batch（批次）的定义
• 为什么要批量处理数据集
• 如何在 Python 中将数据集拆分成批次
• 使用数据集 map() 的 batched 选项的方法

让我们开始吧！

什么是 Batch（批次）？

在机器学习和数据处理领域，批次（batch）无非就是一个数据集的子集。对于大规模数据的处理而言，批次通常能帮助更高效地管理海量数据。与一次性处理整个数据集相比，将数据集拆分成更小的块（也称为批次）可以减少内存消耗并提升计算效率。

例如，假设您有如下的一个 CSV 格式的示例数据：

id,name,age,score
1,John,28,85
2,Jane,34,90
3,Bob,25,72
4,Alice,30,88
5,Charlie,29,91
6,David,35,79
7,Eve,22,95
8,Frank,31,82
9,Grace,27,86
10,Hannah,26,80

上述数据集的一批（batch）示例如下：

6,David,35,79
7,Eve,22,95
8,Frank,31,82
9,Grace,27,86
10,Hannah,26,80

这段批次构成了原始数据集从第 5 行到第 10 行的数据。

批量处理数据集的好处

假设您已经通过某些数据获取技术创建了一个数据集。如果您不熟悉这个过程，可以参考我们关于如何创建数据集的指南。

现在，为什么要将这个数据集进行批量处理？原因在于可以带来以下好处：

• 内存效率：一次只处理更小、更易管理的部分，而不是将整个数据集全部加载到内存中。
• 更快的处理速度：批量处理可以并行进行，从而减少处理大型数据集所需的时间。
• 对 ML 模型更好的训练：可以通过逐步更新权重的方式来训练机器学习模型，从而获得更稳定且更快的收敛效果。
• 更好的可扩展性：使处理大型数据集更容易扩展，即使数据无法一次性全部放入内存。

如何将数据集拆分为批次：五大方法

在深入探讨如何使用 Python 将数据集拆分为批次之前，我们应该先确定一些评估方法的标准。以下是一个有意义的考虑清单：

• 实现方式：用一个代码片段展示如何在一个简单的示例中使用该方法。
• 使用场景：在何种实际情境中适合使用这种拆分数据集的方法。
• 输入类型：该方法支持的数据结构和文件类型。
• 优点：该方法能带来哪些好处。
• 缺点：该方法的限制或不足。

接下来让我们一一分析！

方法 #1：数组切片（Array Slicing）

数组切片是一种将数据集拆分为更小、更易管理批次的简单直接方法。核心思想是将数据集（可以是 list、array 或其他序列）按切片的方式分割成一个个小块。

👨‍💻 实现方式：

def create_batches(data, batch_size):
    return [data[i:i + batch_size] for i in range(0, len(data), batch_size)]

# example usage
data = list(range(1, 51))  # sample dataset
batches = create_batches(data, batch_size=5)

print(batches) 

# output: [[1, 2, 3, 4, 5], ..., [46, 47, 48, 49, 50]]

🎯 使用场景：

• 数据预处理任务，且内存限制不大的场合
• 需要在内存中管理的小批量并行数据处理
• 数据管线中的简单批次处理

🔠 输入类型：

• Python 中的 list、array、tuple
• NumPy 数组
• 将 CSV 数据读取进内存后形成的行列表
• 转换为 list 或 array 的 Pandas DataFrame

👍 优点：

• 实现简单易懂
• 不需要额外的外部库
• 可直接控制批次大小

👎 缺点：

• 受限于可用内存
• 不适合极大型的数据集或复杂的数据结构
• 需要自行实现对数据打乱（shuffling）的逻辑

方法 #2：生成器（Generators）

Python 生成器可以逐批产出数据集。若您不熟悉生成器，它是一种特殊的函数类型，使用 yield 关键字返回迭代器对象。这样，批次就可以在 for 循环或 next() 函数的配合下被顺序访问。

👨‍💻 实现方式：

def data_generator(data, batch_size):
    for i in range(0, len(data), batch_size):
        yield data[i:i + batch_size]

# example usage
data = list(range(1, 51))  # sample dataset
for batch in data_generator(data, batch_size=5):
    print(batch)

# output: 
# [1, 2, 3, 4, 5]
# ...
# [46, 47, 48, 49, 50]

🎯 使用场景：

• 数据管线中的批次处理
• 大规模数据的预处理与数据增强
• 数据管线中从简单到复杂的批次处理

🔠 输入类型：

• Python 中的 list、array、tuple
• NumPy 数组
• 文件型数据集（若可以批量读入）

👍 优点：

• 无需一次性将整个数据加载到内存，可以处理较大的数据集
• 实现简单，配置少
• 可以按需受控地加载数据

👎 缺点：

• 如果需要改变数据顺序，需要另行实现打乱逻辑
• 对动态或可变批次大小的场景支持不够便捷
• 在多线程等并行操作中可能并不是最优解

方法 #3：PyTorch DataLoader

DataLoader 是来自 PyTorch 的一个类，可高效地将数据集拆分成批次。它是一种针对数据集而设计的结构化数据处理方式，还带有诸如打乱数据（shuffling）和并行加载等功能。

需要注意，DataLoader 通常配合 TensorDataset 一起使用。后者是另一个专门用于表示数据集的 PyTorch 数据结构，通常需要传入以下两个参数：

1. inps：输入数据，一般以 Tensor 的形式表现。
2. tgts：对应于输入数据的标签或目标值，也通常用 Tensor 的形式表示。

TensorDataset 将 data 与 target 绑定到一起，以便 DataLoader 在批量和训练时更方便地加载。

👨‍💻 实现方式：

from torch.utils.data import DataLoader, TensorDataset
import torch

# data to define a simple dataset
inputs = torch.arange(1, 51).float().reshape(-1, 1)  # a 1D tensor dataset (input)
targets = inputs ** 2  # square of the input values (simulating a regression task)

# create a TensorDataset and DataLoader
dataset = TensorDataset(inputs, targets)
dataloader = DataLoader(dataset, batch_size=5, shuffle=True)

# iterate through DataLoader
for batch in dataloader:
    print(batch)

# sample output:
# [tensor([[46.],
#         [42.],
#         [25.],
#         [10.],
#         [34.]]), tensor([[2116.],
#         [1764.],
#         [ 625.],
#         [ 100.],
#         [1156.]])]
# ...
# [tensor([[21.],
#         [ 9.],
#         [ 2.],
#         [38.],
#         [44.]]), tensor([[ 441.],
#         [  81.],
#         [   4.],
#         [1444.],
#         [1936.]])] 

🎯 使用场景：

• PyTorch 中的机器学习模型训练与测试
• 为训练批次进行打乱，从而获得相对无偏的数据
• 深度学习任务中对大规模数据的处理

🔠 输入类型：

• PyTorch TensorDataset 格式的自定义数据集
• 以 Tensor 形式表现的表格和数值数据

👍 优点：

• 在进行分批和打乱时对大规模数据集进行优化
• 支持并行数据加载，加快批次检索
• 与 PyTorch 模型和训练循环无缝配合
• 可支持 GPU 加速处理

👎 缺点：

• 需要安装 PyTorch
• 数据需要转换成 Tensor 格式
• 不太适合非机器学习场景的批次处理

方法 #4：TensorFlow batch() 方法

在 TensorFlow 中可以利用 Dataset 的 batch() 方法对数据集进行批次拆分。该方法能将数据集切分成更小的块，并提供并行化、数据处理顺序控制以及命名等功能。

作为一个机器学习库，TensorFlow 还提供了打乱（shuffle）、重复（repeat）和预取（prefetch）等辅助功能。

👨‍💻 实现方式：

import tensorflow as tf

# create a sample dataset
inputs = tf.range(1, 51, dtype=tf.float32)  # a 1D tensor dataset (inputs)
targets = inputs ** 2  # square of input values (simulating a regression task)

# convert the inputs and targets into tf.data.Dataset
inputs_dataset = tf.data.Dataset.from_tensor_slices(inputs)
targets_dataset = tf.data.Dataset.from_tensor_slices(targets)

# create a dataset by zipping the inputs and targets together
dataset = tf.data.Dataset.zip((inputs_dataset, targets_dataset))

# produce a batched dataset
batched_dataset = dataset.batch(batch_size=5)

for batch in batched_dataset:
    print(batch)

# output:
# (<tf.Tensor: shape=(5,), dtype=float32, numpy=array([1., 2., 3., 4., 5.], dtype=float32)>, <tf.Tensor: shape=(5,), dtype=float32, numpy=array([ 1.,  4.,  9., 16., 25.], dtype=float32)>)
# ...
# (<tf.Tensor: shape=(5,), dtype=float32, numpy=array([46., 47., 48., 49., 50.], dtype=float32)>, <tf.Tensor: shape=(5,), dtype=float32, numpy=array([2116., 2209., 2304., 2401., 2500.], dtype=float32)>)

🎯 使用场景：

• 在 TensorFlow 中训练及测试机器学习模型
• 为训练批次进行打乱，从而获得相对无偏的数据
• 深度学习任务中处理大规模数据

🔠 输入类型：

• TensorFlow tf.data.Dataset 对象
• NumPy 数组（可转化为 Dataset）
• TFRecord 文件（TensorFlow 常用的大规模数据二进制格式）

👍 优点：

• 在内存使用方面进行了优化
• 与 TensorFlow 生态系统无缝集成，可用于模型训练与评估
• 支持打乱、预取等实用特性
• 可支持多种数据格式，包括图像、文本和结构化数据

👎 缺点：

• 需要安装 TensorFlow
• 对于更复杂的数据集，需要额外的设置来正确格式化和预处理数据
• 对于小型数据集，批次化可能会引入额外的开销

方法 #5：HDF5 格式

HDF5 是一种被广泛采用的数据格式，适用于管理大型或分层结构的数据。它支持将大型数据集按块进行存储并高效访问。

Python 库 h5py 能够帮助操作 HDF5 文件，并将其加载为 NumPy 数据结构，从而可以实现按需访问具体数据块来进行批处理。

👨‍💻 实现方式：

import h5py
import numpy as np

# load and batch the data from HDF5 file
def load_data_in_batches(batch_size=10):
    # open an HDF5 file
    with h5py.File("dataset.h5", "r") as f:
        inputs = f["input"]
        targets = f["target"]

        # batching the data with an iterator from the disk
        for i in range(0, len(data), batch_size):
            yield inputs[i:i+batch_size], targets[i:i+batch_size]

# iterate through batches
for batch_data, batch_target in load_data_in_batches():
    print("Input Batch:", batch_input)
    print("Targe Batch:", batch_target)

# output:
# Input Batch: [[ 1]
#  [ 2]
#  [ 3]
#  [ 4]
#  [ 5]
#  [ 6]
#  [ 7]
#  [ 8]
#  [ 9]
#  [10]]
# Target Batch: [[  1]
#  [  4]
#  [  9]
#  [ 16]
#  [ 25]
#  [ 36]
#  [ 49]
#  [ 64]
#  [ 81]
#  [100]]
#  ...
# Input Batch: [[41]
#  [42]
#  [43]
#  [44]
#  [45]
#  [46]
#  [47]
#  [48]
#  [49]
#  [50]]
# Target Batch: [[1681]
#  [1764]
#  [1849]
#  [1936]
#  [2025]
#  [2116]
#  [2209]
#  [2304]
#  [2401]
#  [2500]]

🎯 使用场景：

• 非常适合无法全部载入内存的超大型数据集
• 适合处理多维度数据
• 非常适合在机器学习任务中以高效、可压缩的方式在磁盘中存储和获取数据

🔠 输入类型：

• HDF5 文件

👍 优点：

• HDF5 支持数据压缩和分块，大幅减少存储需求
• 可以对大型数据集的部分内容进行随机访问，而无需加载全部
• 可以在单个文件中存储多个数据集，适合管理复杂数据
• 被众多科学库支持，包括 NumPy、TensorFlow 和 PyTorch

👎 缺点：

• 需要额外学习和配置 HDF5 格式
• 依赖 h5py 库来完整操作 HDF5 文件
• 并非所有数据集都提供 HDF5 版本

其他方案

除了以上几种将数据集拆分成批次的方法，类似方案还有不少。

另一个可行的方案是使用 Hugging Face 的 datasets 库。它提供了对整个数据集进行转换并进行批量处理的功能。通过设置 batched=True，您就可以在不手动切分数据的情况下直接对数据集进行批量级别的转换，比如：

from datasets import load_dataset

# load a sample dataset
dataset = load_dataset("imdb", split="train")

# define a batch processing function
def process_batch(batch):
    # simple tokenization task
    return {"tokens": [text.split() for text in batch["text"]]}

# apply batch processing
batched_dataset = dataset.map(process_batch, batched=True, batch_size=32)

这里的 dataset map() batched=True 参数特别适用于需要批量级别转换（如分词）且要减少内存占用、加速数据处理的场景。对于 NLP 以及机器学习中的文本和表格数据处理来说，这种方式尤为高效。

结论

在这篇关于如何将数据集拆分成批次的指南中，我们探讨了在 Python 中进行数据切分的最佳方法、库与方案，目的是把大型数据分割成更易管理的部分，从而简化并加速数据处理。

无论您选择哪种方法，所有方案的前提都是要先获取到包含所需数据的数据集。不过，学术上开源的数据集并不能涵盖所有情境。

如果您需要从金融到电影等各种类别的数据集，不妨查看Bright Data 的数据集商店 (Dataset Marketplace)。其中包含了从各大网站抓取的数百个数据集，并根据以下主题进行分类：

• 商业数据集：来自 LinkedIn、CrunchBase、Owler、Indeed 等关键来源的数据。
• 电商数据集：来自 Amazon、Walmart、Target、Zara、Zalando、Asos 等的电商数据。
• 房产数据集：来自 Zillow、MLS 等网站的房产数据。
• 社交媒体数据集：来自 Facebook、Instagram、YouTube、Reddit 等平台的社交媒体数据。
• 金融数据集：来自 Yahoo Finance、Market Watch、Investopedia 等站点的数据。

如果这些已做好的数据集仍然不能满足您的需求，也可以考虑使用我们的自定义数据采集服务。

此外，Bright Data 还提供了多种强力的抓取工具，如 Web Scraper APIs 以及 Scraping Browser。

现在即可注册一个 Bright Data 免费账号，开始探索这些数据集吧！