MegEngine 使用小技巧：借助 DataLoader 获取分批数据

在使用 MegEngine 进行模型训练时，首先要进行的是数据加载和预处理。在此过程中，MegEngine 中的​​megengine.data​​模块，提供了数据分批功能，其内部实现流程如下图：

通过使用 Dataloader，我们可以方便地对数据进行相关操作，比如设置 batch_size 以及自定义数据预处理的过程。示例代码如下：

from megengine.data import DataLoader
from megengine.data.dataset import ArrayDataset
from megengine.data.sampler import SequentialSampler

def init_dataset():
    rand_data = np.random.randint(0, 255, size=(100, 1, 32, 32), dtype=np.uint8)
    label = np.random.randint(0, 10, size=(100,), dtype=int)
    dataset = ArrayDataset(rand_data, label)
    return dataset

dataset = init_dataset()
dataloader = DataLoader(
    dataset, sampler=SequentialSampler(dataset, batch_size=5)
)

for batch_data, batch_label in dataloader:
    print(batch_data.shape)
    break

在上面的代码中，我们用 ​​ArrayDataset​​ 对 NumPy ndarray 格式的数据集进行了快速封装， 接着使用顺序采样器 ​​SequentialSampler​​ 对 house_dataset 进行了采样， 二者用来作为参数初始化 ​​DataLoader​​, 最终获取到了一个可迭代的对象，每次提供 batch_size 大小的数据和标记。

len(dataloader)    # 20
batch_data.shape   # (5, 1, 32, 32)
batch_label.shape  # (5,)

我们在上面选定的 batch_size 为 5，样本容量为 100, 因此可以划分成 20 批数据。

Dataloader 也支持用户自定义对数据做预处理操作，而在 MegEngine 中也定义很多图像处理领域的预处理方式，如图像裁剪、翻转等。 更多 DataLoader 功能介绍，见文档：<​​https://www.megengine.org.cn/doc/stable/zh/user-guide/model-development/data/index.html​​>

附

「MegEngine 使用小技巧」系列文章，重点输出 MegEngine 及周边工具的使用技巧，如有催更或投稿，欢迎联系我们哦～