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

原创
2023/03/09 10:31
阅读数 86

在使用 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 及周边工具的使用技巧,如有催更或投稿,欢迎联系我们哦~

技术交流 QQ 群:1029741705;Bot 微信:megengine-bot

更多 MegEngine 信息获取,您可以:查看文档GitHub 项目,欢迎参与 MegEngine 社区贡献,成为 Awesome MegEngineer,荣誉证书、定制礼品享不停。

展开阅读全文
加载中
点击引领话题📣 发布并加入讨论🔥
0 评论
0 收藏
0
分享
返回顶部
顶部