# 与 NumPy 配合使用

本文档简要介绍了如何使用 `datasets` 与 NumPy 配合，重点说明如何从数据集中获取 `numpy.ndarray` 对象，以及如何利用它们来训练基于 NumPy 的模型，例如 `scikit-learn` 模型。

## 数据集格式

默认情况下，数据集返回的是普通的 Python 对象：整数、浮点数、字符串、列表等。

若要获取 NumPy 数组，可以将数据集的格式设置为 `numpy`：

```py
>>> from datasets import Dataset
>>> data = [[1, 2], [3, 4]]
>>> ds = Dataset.from_dict({"data": data})
>>> ds = ds.with_format("numpy")
>>> ds[0]
{'data': array([1, 2])}
>>> ds[:2]
{'data': array([
    [1, 2],
    [3, 4]])}
```

> [!TIP]
> [Dataset](/docs/datasets/v5.2.0/zh/package_reference/main_classes#datasets.Dataset) 对象是 Arrow 表的封装，它允许从数据集的数组快速读取到 NumPy 数组。

请注意，完全相同的操作也适用于 `DatasetDict` 对象，因此当将 `DatasetDict` 的格式设置为 `numpy` 时，其中的所有 `Dataset` 都会被格式化为 `numpy`：

```py
>>> from datasets import DatasetDict
>>> data = {"train": {"data": [[1, 2], [3, 4]]}, "test": {"data": [[5, 6], [7, 8]]}}
>>> dds = DatasetDict.from_dict(data)
>>> dds = dds.with_format("numpy")
>>> dds["train"][:2]
{'data': array([
    [1, 2],
    [3, 4]])}
```

### N 维数组

如果您的数据集由 N 维数组组成，您会发现默认情况下，如果形状固定，它们会被视为同一个数组：

```py
>>> from datasets import Dataset
>>> data = [[[1, 2],[3, 4]], [[5, 6],[7, 8]]]  # fixed shape
>>> ds = Dataset.from_dict({"data": data})
>>> ds = ds.with_format("numpy")
>>> ds[0]
{'data': array([[1, 2],
        [3, 4]])}
```

```py
>>> from datasets import Dataset
>>> data = [[[1, 2],[3]], [[4, 5, 6],[7, 8]]]  # varying shape
>>> ds = Dataset.from_dict({"data": data})
>>> ds = ds.with_format("numpy")
>>> ds[0]
{'data': array([array([1, 2]), array([3])], dtype=object)}
```

然而，这种逻辑通常需要缓慢的形状比较和数据复制。为避免这种情况，您必须显式使用 `Array` 特征类型，并指定张量的形状：

```py
>>> from datasets import Dataset, Features, Array2D
>>> data = [[[1, 2],[3, 4]],[[5, 6],[7, 8]]]
>>> features = Features({"data": Array2D(shape=(2, 2), dtype='int32')})
>>> ds = Dataset.from_dict({"data": data}, features=features)
>>> ds = ds.with_format("numpy")
>>> ds[0]
{'data': array([[1, 2],
        [3, 4]])}
>>> ds[:2]
{'data': array([[[1, 2],
         [3, 4]],
 
        [[5, 6],
         [7, 8]]])}
```

### 其他特征类型

[ClassLabel](/docs/datasets/v5.2.0/zh/package_reference/main_classes#datasets.ClassLabel) 数据被正确转换为数组：

```py
>>> from datasets import Dataset, Features, ClassLabel
>>> labels = [0, 0, 1]
>>> features = Features({"label": ClassLabel(names=["negative", "positive"])})
>>> ds = Dataset.from_dict({"label": labels}, features=features)
>>> ds = ds.with_format("numpy")
>>> ds[:3]
{'label': array([0, 0, 1])}
```

字符串和二进制对象保持不变，因为 NumPy 仅支持数字。

[Image](/docs/datasets/v5.2.0/zh/package_reference/main_classes#datasets.Image) 和 [Audio](/docs/datasets/v5.2.0/zh/package_reference/main_classes#datasets.Audio) 特征类型也受支持。

> [!TIP]
> 要使用 [Image](/docs/datasets/v5.2.0/zh/package_reference/main_classes#datasets.Image) 功能类型，您需要安装 `vision` 附加组件，命令为 `pip install datasets[vision]`。

```py
>>> from datasets import Dataset, Features, Image
>>> images = ["path/to/image.png"] * 10
>>> features = Features({"image": Image()})
>>> ds = Dataset.from_dict({"image": images}, features=features)
>>> ds = ds.with_format("numpy")
>>> ds[0]["image"].shape
(512, 512, 3)
>>> ds[0]
{'image': array([[[ 255, 255, 255],
              [ 255, 255, 255],
              ...,
              [ 255, 255, 255],
              [ 255, 255, 255]]], dtype=uint8)}
>>> ds[:2]["image"].shape
(2, 512, 512, 3)
>>> ds[:2]
{'image': array([[[[ 255, 255, 255],
              [ 255, 255, 255],
              ...,
              [ 255, 255, 255],
              [ 255, 255, 255]]]], dtype=uint8)}
```

> [!TIP]
> 要使用 [Audio](/docs/datasets/v5.2.0/zh/package_reference/main_classes#datasets.Audio) 功能类型，您需要安装 `audio` 附加组件，命令为 `pip install datasets[audio]`。

```py
>>> from datasets import Dataset, Features, Audio
>>> audio = ["path/to/audio.wav"] * 10
>>> features = Features({"audio": Audio()})
>>> ds = Dataset.from_dict({"audio": audio}, features=features)
>>> ds = ds.with_format("numpy")
>>> ds[0]["audio"]["array"]
array([-0.059021  , -0.03894043, -0.00735474, ...,  0.0133667 ,
              0.01809692,  0.00268555], dtype=float32)
>>> ds[0]["audio"]["sampling_rate"]
array(44100, weak_type=True)
```

## 数据加载

NumPy 没有内置的数据加载功能，因此你需要将 `X, y` 这样的 NumPy 数组实例化以便在 `scikit-learn` 中使用，或者使用 [PyTorch](https://pytorch.org/) 等库通过 `DataLoader` 来加载数据。

### 使用 `with_format('numpy')`

从数据集获取 NumPy 数组最简单的方法是使用 `with_format('numpy')` 方法。假设我们想在 HuggingFace Hub 上可用的 [MNIST 数据集](http://yann.lecun.com/exdb/mnist/)（https://huggingface.co/datasets/mnist）上训练一个神经网络。

```py
>>> from datasets import load_dataset
>>> ds = load_dataset("ylecun/mnist")
>>> ds = ds.with_format("numpy")
>>> ds["train"][0]
{'image': array([[  0,   0,   0, ...],
                       [  0,   0,   0, ...],
                       ...,
                       [  0,   0,   0, ...],
                       [  0,   0,   0, ...]], dtype=uint8),
 'label': array(5)}
```

设置好格式后，我们可以使用 `Dataset.iter()` 方法，以 NumPy 为基础，将数据集分批输入模型：

```py
>>> for epoch in range(epochs):
...     for batch in ds["train"].iter(batch_size=32):
...         x, y = batch["image"], batch["label"]
...         ...
```