199 lines
12 KiB
Markdown
199 lines
12 KiB
Markdown
# 数据集（Dataset）与数据加载器（DataLoader）
|
||
|
||
```{hint}
|
||
如果你没有接触过 PyTorch 的数据集与数据加载器，我们推荐先浏览 [PyTorch 官方教程](https://pytorch.org/tutorials/beginner/basics/data_tutorial.html)以了解一些基本概念
|
||
```
|
||
|
||
数据集与数据加载器是 MMEngine 中训练流程的必要组件，它们的概念来源于 PyTorch，并且在含义上与 PyTorch 保持一致。通常来说，数据集定义了数据的总体数量、读取方式以及预处理，而数据加载器则在不同的设置下迭代地加载数据，如批次大小（`batch_size`）、随机乱序（`shuffle`）、并行（`num_workers`）等。数据集经过数据加载器封装后构成了数据源。在本篇教程中，我们将按照从外（数据加载器）到内（数据集）的顺序，逐步介绍它们在 MMEngine 执行器中的用法，并给出一些常用示例。读完本篇教程，你将会：
|
||
|
||
- 掌握如何在 MMEngine 的执行器中配置数据加载器
|
||
- 学会在配置文件中使用已有（如 `torchvision`）数据集
|
||
- 了解如何使用自己的数据集
|
||
|
||
## 数据加载器详解
|
||
|
||
在执行器（`Runner`）中，你可以分别配置以下 3 个参数来指定对应的数据加载器
|
||
|
||
- `train_dataloader`：在 `Runner.train()` 中被使用，为模型提供训练数据
|
||
- `val_dataloader`：在 `Runner.val()` 中被使用，也会在 `Runner.train()` 中每间隔一段时间被使用，用于模型的验证评测
|
||
- `test_dataloader`：在 `Runner.test()` 中被使用，用于模型的测试
|
||
|
||
MMEngine 完全支持 PyTorch 的原生 `DataLoader`，因此上述 3 个参数均可以直接传入构建好的 `DataLoader`，如[15分钟上手](../get_started/15_minutes.md)中的例子所示。同时，借助 MMEngine 的[注册机制](../advanced_tutorials/registry.md)，以上参数也可以传入 `dict`，如下面代码（以下简称例 1）所示。字典中的键值与 `DataLoader` 的构造参数一一对应。
|
||
|
||
```python
|
||
runner = Runner(
|
||
train_dataloader=dict(
|
||
batch_size=32,
|
||
sampler=dict(
|
||
type='DefaultSampler',
|
||
shuffle=True),
|
||
dataset=torchvision.datasets.CIFAR10(...),
|
||
collate_fn=dict(type='default_collate')
|
||
)
|
||
)
|
||
```
|
||
|
||
在这种情况下，数据加载器会在实际被用到时，在执行器内部被构建。
|
||
|
||
```{note}
|
||
关于 `DataLoader` 的更多可配置参数，你可以参考 [PyTorch API 文档](https://pytorch.org/docs/stable/data.html#torch.utils.data.DataLoader)
|
||
```
|
||
|
||
```{note}
|
||
如果你对于构建的具体细节感兴趣，你可以参考 [build_dataloader](mmengine.runner.Runner.build_dataloader)
|
||
```
|
||
|
||
细心的你可能会发现，例 1 **并非**直接由[15分钟上手](../get_started/15_minutes.md)中的示例代码简单修改而来。你可能本以为将 `DataLoader` 简单替换为 `dict` 就可以无缝切换，但遗憾的是，基于注册机制构建时 MMEngine 会有一些隐式的转换和约定。我们将介绍其中的不同点，以避免你使用配置文件时产生不必要的疑惑。
|
||
|
||
### sampler 与 shuffle
|
||
|
||
与 15 分钟上手明显不同，例 1 中我们添加了 `sampler` 参数，这是由于在 MMEngine 中我们要求通过 `dict` 传入的数据加载器的配置**必须包含 `sampler` 参数**。同时，`shuffle` 参数也从 `DataLoader` 中移除，这是由于在 PyTorch 中 **`sampler` 与 `shuffle` 参数是互斥的**，见 [PyTorch API 文档](https://pytorch.org/docs/stable/data.html#torch.utils.data.DataLoader)。
|
||
|
||
```{note}
|
||
事实上，在 PyTorch 的实现中，`shuffle` 只是一个便利记号。当设置为 `True` 时 `DataLoader` 会自动在内部使用 `RandomSampler`
|
||
```
|
||
|
||
当考虑 `sampler` 时，例 1 代码**基本**可以认为等价于下面的代码块
|
||
|
||
```python
|
||
from mmengine.dataset import DefaultSampler
|
||
|
||
dataset = torchvision.datasets.CIFAR10(...)
|
||
sampler = DefaultSampler(dataset, shuffle=True)
|
||
|
||
runner = Runner(
|
||
train_dataloader=DataLoader(
|
||
batch_size=32,
|
||
sampler=sampler,
|
||
dataset=dataset,
|
||
collate_fn=default_collate
|
||
)
|
||
)
|
||
```
|
||
|
||
```{warning}
|
||
上述代码的等价性只有在：1）使用单进程训练，以及 2）没有配置执行器的 `randomness` 参数时成立。这是由于使用 `dict` 传入 `sampler` 时，执行器会保证它在分布式训练环境设置完成后才被惰性构造，并接收到正确的随机种子。这两点在手动构造时需要额外工作且极易出错。因此，上述的写法只是一个示意而非推荐写法。我们**强烈建议 `sampler` 以 `dict` 的形式传入**，让执行器处理构造顺序，以避免出现问题。
|
||
```
|
||
|
||
### DefaultSampler
|
||
|
||
上面例子可能会让你好奇：`DefaultSampler` 是什么，为什么要使用它，是否有其他选项？事实上，`DefaultSampler` 是 MMEngine 内置的一种采样器，它屏蔽了单进程训练与多进程训练的细节差异，使得单卡与多卡训练可以无缝切换。如果你有过使用 PyTorch `DistributedDataParallel` 的经验，你一定会对其中更换数据加载器的 `sampler` 参数有所印象。但在 MMEngine 中，这一细节通过 `DefaultSampler` 而被屏蔽。
|
||
|
||
除了 `Dataset` 本身之外，`DefaultSampler` 还支持以下参数配置：
|
||
|
||
- `shuffle` 设置为 `True` 时会打乱数据集的读取顺序
|
||
- `seed` 打乱数据集所用的随机种子，通常不需要在此手动设置，会从 `Runner` 的 `randomness` 入参中读取
|
||
- `round_up` 设置为 `True` 时，与 PyTorch `DataLoader` 中设置 `drop_last=False` 行为一致。如果你在迁移 PyTorch 的项目，你可能需要注意这一点。
|
||
|
||
```{note}
|
||
更多关于 `DefaultSampler` 的内容可以参考 [API 文档](mmengine.dataset.DefaultSampler)
|
||
```
|
||
|
||
`DefaultSampler` 适用于绝大部分情况，并且我们保证在执行器中使用它时，随机数等容易出错的细节都被正确地处理，防止你陷入多进程训练的常见陷阱。如果你想要使用基于迭代次数 (iteration-based) 的训练流程，你也许会对 [InfiniteSampler](mmengine.dataset.InfiniteSampler) 感兴趣。如果你有更多的进阶需求，你可能会想要参考上述两个内置 `sampler` 的代码，实现一个自定义的 `sampler` 并注册到 `DATA_SAMPLERS` 根注册器中。
|
||
|
||
```python
|
||
@DATA_SAMPLERS.register_module()
|
||
class MySampler(Sampler):
|
||
pass
|
||
|
||
runner = Runner(
|
||
train_dataloader=dict(
|
||
sampler=dict(type='MySampler'),
|
||
...
|
||
)
|
||
)
|
||
```
|
||
|
||
### 不起眼的 collate_fn
|
||
|
||
PyTorch 的 `DataLoader` 中，`collate_fn` 这一参数常常被使用者忽略，但在 MMEngine 中你需要额外注意：当你传入 `dict` 来构造数据加载器时，MMEngine 会默认使用内置的 [pseudo_collate](mmengine.dataset.pseudo_collate)，这一点明显区别于 PyTorch 默认的 [default_collate](https://pytorch.org/docs/stable/data.html#torch.utils.data.default_collate)。因此，当你迁移 PyTorch 项目时，需要在配置文件中手动指明 `collate_fn` 以保持行为一致。
|
||
|
||
```{note}
|
||
MMEngine 中使用 `pseudo_collate` 作为默认值，主要是由于历史兼容性原因，你可以不必过于深究，只需了解并避免错误使用即可。
|
||
```
|
||
|
||
MMengine 中提供了 2 种内置的 `collate_fn`：
|
||
|
||
- `pseudo_collate`，缺省时的默认参数。它不会将数据沿着 `batch` 的维度合并。详细说明可以参考 [pseudo_collate](mmengine.dataset.pseudo_collate)
|
||
- `default_collate`，与 PyTorch 中的 `default_collate` 行为几乎完全一致，会将数据转化为 `Tensor` 并沿着 `batch` 维度合并。一些细微不同和详细说明可以参考 [default_collate](mmengine.dataset.default_collate)
|
||
|
||
如果你想要使用自定义的 `collate_fn`，你也可以将它注册到 `FUNCTIONS` 根注册器中来使用
|
||
|
||
```python
|
||
@FUNCTIONS.register_module()
|
||
def my_collate_func(data_batch: Sequence) -> Any:
|
||
pass
|
||
|
||
runner = Runner(
|
||
train_dataloader=dict(
|
||
...
|
||
collate_fn=dict(type='my_collate_func')
|
||
)
|
||
)
|
||
```
|
||
|
||
## 数据集详解
|
||
|
||
数据集通常定义了数据的数量、读取方式与预处理，并作为参数传递给数据加载器供后者分批次加载。由于我们使用了 PyTorch 的 `DataLoader`，因此数据集也自然与 PyTorch `Dataset` 完全兼容。同时得益于注册机制，当数据加载器使用 `dict` 在执行器内部构建时，`dataset` 参数也可以使用 `dict` 传入并在内部被构建。这一点使得编写配置文件成为可能。
|
||
|
||
### 使用 torchvision 数据集
|
||
|
||
`torchvision` 中提供了丰富的公开数据集，它们都可以在 MMEngine 中直接使用，例如 [15 分钟上手](../get_started/15_minutes.md)中的示例代码就使用了其中的 `Cifar10` 数据集，并且使用了 `torchvision` 中内置的数据预处理模块。
|
||
|
||
但是，当需要将上述示例转换为配置文件时，你需要对 `torchvision` 中的数据集进行额外的注册。如果你同时用到了 `torchvision` 中的数据预处理模块，那么你也需要编写额外代码来对它们进行注册和构建。下面我们将给出一个等效的例子来展示如何做到这一点。
|
||
|
||
```python
|
||
import torchvision.transforms as tvt
|
||
from mmengine.registry import DATASETS, TRANSFORMS
|
||
from mmengine.dataset.base_dataset import Compose
|
||
|
||
# 注册 torchvision 的 CIFAR10 数据集
|
||
# 数据预处理也需要在此一起构建
|
||
@DATASETS.register_module(name='Cifar10', force=False)
|
||
def build_torchvision_cifar10(transform=None, **kwargs):
|
||
if isinstance(transform, dict):
|
||
transform = [transform]
|
||
if isinstance(transform, (list, tuple)):
|
||
transform = Compose(transform)
|
||
return torchvision.datasets.CIFAR10(**kwargs, transform=transform)
|
||
|
||
# 注册 torchvision 中用到的数据预处理模块
|
||
DATA_TRANSFORMS.register_module('RandomCrop', module=tvt.RandomCrop)
|
||
DATA_TRANSFORMS.register_module('RandomHorizontalFlip', module=tvt.RandomHorizontalFlip)
|
||
DATA_TRANSFORMS.register_module('ToTensor', module=tvt.ToTensor)
|
||
DATA_TRANSFORMS.register_module('Normalize', module=tvt.Normalize)
|
||
|
||
# 在 Runner 中使用
|
||
runner = Runner(
|
||
train_dataloader=dict(
|
||
batch_size=32,
|
||
sampler=dict(
|
||
type='DefaultSampler',
|
||
shuffle=True),
|
||
dataset=dict(type='Cifar10',
|
||
root='data/cifar10',
|
||
train=True,
|
||
download=True,
|
||
transform=[
|
||
dict(type='RandomCrop', size=32, padding=4),
|
||
dict(type='RandomHorizontalFlip'),
|
||
dict(type='ToTensor'),
|
||
dict(type='Normalize', **norm_cfg)])
|
||
)
|
||
)
|
||
```
|
||
|
||
```{note}
|
||
上述例子中大量使用了[注册机制](../advanced_tutorials/registry.md)，并且用到了 MMEngine 中的 [Compose](mmengine.dataset.Compose)。如果你急需在配置文件中使用 `torchvision` 数据集，你可以参考上述代码并略作修改。但我们更加推荐你有需要时在下游库（如 [MMDet](https://github.com/open-mmlab/mmdetection) 和 [MMCls](https://github.com/open-mmlab/mmclassification) 等）中寻找对应的数据集实现，从而获得更好的使用体验。
|
||
```
|
||
|
||
### 自定义数据集
|
||
|
||
你可以像使用 PyTorch 一样，自由地定义自己的数据集，或将之前 PyTorch 项目中的数据集拷贝过来。如果你想要了解如何自定义数据集，可以参考 [PyTorch 官方教程](https://pytorch.org/tutorials/beginner/basics/data_tutorial.html#creating-a-custom-dataset-for-your-files)
|
||
|
||
### 使用 MMEngine 的数据集基类
|
||
|
||
除了直接使用 PyTorch 的 `Dataset` 来自定义数据集之外，你也可以使用 MMEngine 内置的 `BaseDataset`，参考[数据集基类](../advanced_tutorials/basedataset.md)文档。它对标注文件的格式做了一些约定，使得数据接口更加统一、多任务训练更加便捷。同时，数据集基类也可以轻松地搭配内置的[数据变换](../advanced_tutorials/data_transform.md)使用，减轻你从头搭建训练流程的工作量。
|
||
|
||
目前，`BaseDataset` 已经在 OpenMMLab 2.0 系列的下游仓库中被广泛使用。
|