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

在使用MegEngine进行模型训练时，首先要进行的是数据加载和预处理。在此过程中，MegEngine中的megengine.data（https://www.megengine.org.cn/doc/stable/zh/reference/data.html#module-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 （https://www.megengine.org.cn/doc/stable/zh/reference/api/megengine.data.dataset.ArrayDataset.html#megengine.data.dataset.ArrayDataset）对NumPy ndarray格式的数据集进行了快速封装， 接着使用顺序采样器 SequentialSampler （https://www.megengine.org.cn/doc/stable/zh/reference/api/megengine.data.SequentialSampler.html#megengine.data.SequentialSampler）对 house_dataset 进行了采样， 二者用来作为参数初始化 DataLoader（https://www.megengine.org.cn/doc/stable/zh/reference/api/megengine.data.DataLoader.html#megengine.data.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及周边工具的使用技巧，如有催更或投稿，欢迎联系我们哦～

技术交流QQ 群：1029741705；Bot 微信：megengine-bot

更多MegEngine信息获取，您可以：查看文档：

（http://www.megengine.org.cn/doc/stable/zh/）和GitHub项目：

（https://github.com/MegEngine），欢迎参与MegEngine社区贡献，成为Awesome MegEngineer（http://www.megengine.org.cn/community-AMGE），荣誉证书、定制礼品享不停。